Update vid2persona/init.py

vid2persona/init.py

@@ -20,7 +20,7 @@ def init_model(model_id):
     
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     tokenizer.use_default_system_prompt = False
-    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="cuda")
 
 def auth_gcp():
     gcp_credentials = os.getenv("GCP_CREDENTIALS")