test device_map auto

methods/utilities.py

@@ -30,10 +30,11 @@ def load_llama_llm(AUTH_TOKEN):
         token=AUTH_TOKEN
     )
     model = AutoModelForCausalLM.from_pretrained(
-            model_id, 
-            #torch_dtype=torch.float16, 
-            trust_remote_code=True,
-            token=AUTH_TOKEN
+        model_id, 
+        torch_dtype=torch.float16,
+        device_map='auto',
+        trust_remote_code=True,
+        token=AUTH_TOKEN
     )
     model = model.to('cuda')
     model = model.eval()