Update app.py

app.py

@@ -5,21 +5,20 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 # Set seed for reproducibility
 torch.random.manual_seed(0)
 
-# Load the model and tokenizer (using CPU)
+# Load the model and tokenizer (using Accelerate)
 model = AutoModelForCausalLM.from_pretrained(
     "microsoft/Phi-3.5-mini-instruct",
-    device_map="cpu",  # Use CPU
+    device_map="auto",  # Managed by Accelerate
     torch_dtype="auto",
     trust_remote_code=True,
 )
 tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3.5-mini-instruct")
 
-# Define the pipeline
+# Define the pipeline (no device argument)
 pipe = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    device=-1  # CPU
 )
 
 # System message (invisible to the user)