Update main.py

main.py

@@ -43,18 +43,13 @@ processor = AutoProcessor.from_pretrained(
     "microsoft/Florence-2-base", trust_remote_code=True
 )
 
-try:
-    model = AutoModelForCausalLM.from_pretrained(
-        "weights/icon_caption_florence",
-        torch_dtype=torch.float16,
-        trust_remote_code=True,
-    ).to("cuda")
-except:
-    model = AutoModelForCausalLM.from_pretrained(
-        "weights/icon_caption_florence",
-        torch_dtype=torch.float16,
-        trust_remote_code=True,
-    )
+
+model = AutoModelForCausalLM.from_pretrained(
+    "weights/icon_caption_florence",
+    torch_dtype=torch.float16,
+    trust_remote_code=True,
+    device_map='auto',
+).to("cuda")
 caption_model_processor = {"processor": processor, "model": model}
 print("finish loading model!!!")