Update app.py

app.py

@@ -101,13 +101,14 @@ def loadModel():
                     device_map="auto",
                     quantization_config=quantization_config,
                     max_memory={
-                        0: "10GB",  # GPU 0 with 20GB memory
-                        1: "22GB",  # GPU 0 with 20GB memory
-                        2: "22GB",  # GPU 0 with 20GB memory
-                        3: "22GB",  # GPU 0 with 20GB 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
                         #"cpu": "5GB"  # CPU with 100GB memory
                     },
                 )
+        print_resources()
         model = PeftModel.from_pretrained(
                     model,
                     lora_weights,
@@ -116,10 +117,10 @@ def loadModel():
                     torch_dtype=torch.float16,
                     is_trainable=False,
                     max_memory={
-                        0: "10GB",  # GPU 0 with 20GB memory
-                        1: "22GB",  # GPU 0 with 20GB memory
-                        2: "22GB",  # GPU 0 with 20GB memory
-                        3: "22GB",  # GPU 0 with 20GB memory
+                        #0: "1GB",  # GPU 0 with 20GB memory
+                        #1: "5GB",  # GPU 0 with 20GB memory
+                        2: "35GB",  # GPU 0 with 20GB memory
+                        3: "40GB",  # GPU 0 with 20GB memory
                         #"cpu": "5GB"  # CPU with 100GB memory
                     },
                 )