Update colab.py

Beyondo · Dec 18, 2022 · f7af7b8 · f7af7b8
1 parent 9f31095
commit f7af7b8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/colab.py b/colab.py
@@ -98,7 +98,7 @@ def init(ModelName, debug=False):
         try:
             install_vendor()
             print("Initializing model " + model_name + ":")
-            os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb=512"
+            os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb=256"
             torch.set_default_dtype(torch.float16)
             rev = "diffusers-115k" if model_name == "naclbit/trinart_stable_diffusion_v2" else "" if model_name == "prompthero/openjourney" else "fp16"
             # Hook VOIDPipeline to StableDiffusionPipeline