Update app.py

app.py

@@ -18,16 +18,16 @@ model = AutoModelForImageTextToText.from_pretrained(
     model_id,
     # load_in_4bit=True,
     torch_dtype=torch.float32 if device.type == "cpu" else torch.bfloat16,
-    device_map="auto" if device.type == "cuda" else None,
+    device_map=device,
 ).to(device)
 
-if device.type == "cuda":
-    model.gradient_checkpointing_enable()
-# model.gradient_checkpointing_enable()
+# if device.type == "cuda":
+#     model.gradient_checkpointing_enable()
+model.gradient_checkpointing_enable()
 
 processor = AutoProcessor.from_pretrained(model_id)
 
-@spaces.GPU(duration=120)
+# @spaces.GPU(duration=120)
 # Function to process the image and generate the description
 def generate_description(image: Image.Image, instruction: str):