Update app.py

app.py

@@ -102,9 +102,9 @@ def loadModel():
                     quantization_config=quantization_config,
                     max_memory={
                         0: "20GB",  # GPU 0 with 20GB memory
-                        1: "45GB",  # GPU 0 with 20GB memory
-                        2: "5GB",  # GPU 0 with 20GB memory
-                        #3: "1GB",  # GPU 0 with 20GB memory
+                        1: "40GB",  # GPU 0 with 20GB memory
+                        2: "40GB",  # GPU 0 with 20GB memory
+                        3: "40GB",  # GPU 0 with 20GB memory
                         #"cpu": "5GB"  # CPU with 100GB memory
                     },
                 )
@@ -117,9 +117,9 @@ def loadModel():
                     torch_dtype=torch.float16,
                     is_trainable=False,
                     max_memory={
-                        #0: "1GB",  # GPU 0 with 20GB memory
-                        #1: "5GB",  # GPU 0 with 20GB memory
-                        2: "35GB",  # GPU 0 with 20GB memory
+                        0: "20GB",  # GPU 0 with 20GB memory
+                        1: "8GB",  # GPU 0 with 20GB memory
+                        2: "8GB",  # GPU 0 with 20GB memory
                         3: "40GB",  # GPU 0 with 20GB memory
                         #"cpu": "5GB"  # CPU with 100GB memory
                     },