Upload model_utils.py

src/model_utils.py

@@ -12,11 +12,9 @@ def load_hf_model(model_name, device="cpu"):
         except NotImplementedError:
             # If meta tensor error occurs, use to_empty()
             model = model.to_empty(device=device)
-        device_id = 0
-    else:
-        device_id = -1
     
-    return pipeline("text-generation", model=model, tokenizer=tokenizer, device=device_id)
+    # Don't specify device in pipeline when using accelerate
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 def generate_answer(text_gen, question, context):
     prompt = f"Context: {context}\n\nQuestion: {question}\n\nAnswer:"