Update app.py

app.py

@@ -26,16 +26,24 @@ model = None
 gc.collect()
 
 logger.info("start")
+has_cuda = torch.cuda.is_available()
 
-model = AutoModelForCausalLM.from_pretrained(
-    "model",   # loc
-    device_map="auto", 
-    torch_dtype=torch.bfloat16, 
-    load_in_8bit=True,
-    trust_remote_code=True,
-    # use_ram_optimized_load=False,
-    # offload_folder="offload_folder",
-)
+if has_cuda:
+    model = AutoModelForCausalLM.from_pretrained(
+        "model",   # loc
+        # device_map="auto", 
+        torch_dtype=torch.bfloat16, 
+        load_in_8bit=True,
+        trust_remote_code=True,
+        # use_ram_optimized_load=False,
+        # offload_folder="offload_folder",
+    ).cuda()
+else:
+    model = (
+        AutoModel.from_pretrained(model_name, trust_remote_code=True).float()
+    ).float()
+
+model = model.eval()
 
 rich.print(f"{model=}")