Update app.py

app.py

@@ -10,7 +10,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     torch_dtype=torch.float32,
-    device="cuda" if torch.cuda.is_available() else "cpu"
+    # device_map="cuda" if torch.cuda.is_available() else "cpu"
   )
 
 llama3_am = pipeline(
@@ -18,7 +18,7 @@ llama3_am = pipeline(
     model=model,
     tokenizer=tokenizer,
     eos_token_id=tokenizer.eos_token_id,
-    device="cuda" if torch.cuda.is_available() else "cpu"
+    # device_map="cuda" if torch.cuda.is_available() else "cpu"
   )
 
 # Function that accepts a prompt and generates text