Fix model loading with proper torch_dtype and device_map

app.py

@@ -26,8 +26,13 @@ MAX_LENGTH = int(os.environ.get("STARVECTOR_MAX_LENGTH", "4000"))
 
 print(f"Starting StarVector Space on device={DEVICE} dtype={DTYPE} model={MODEL_NAME}", flush=True)
 
-starvector = StarVectorForCausalLM.from_pretrained(MODEL_NAME)
-starvector = starvector.to(device=DEVICE)
+starvector = StarVectorForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=DTYPE,
+    device_map="auto" if DEVICE == "cuda" else None
+)
+if DEVICE != "cuda":
+    starvector = starvector.to(device=DEVICE)
 starvector.eval()