fix: device_map to auto

app.py

@@ -3,11 +3,10 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
-device = 'cuda'
-torch_dtype = torch.bfloat16
+device_map = 'auto'
 
 def load_model() -> AutoModelForCausalLM:
-    return AutoModelForCausalLM.from_pretrained(model_name, device=device, torch_dtype=torch_dtype)
+    return AutoModelForCausalLM.from_pretrained(model_name, device_map=device_map)
 
 def load_tokenizer() -> AutoTokenizer:
     return AutoTokenizer.from_pretrained(model_name)