Adds device_map to force GPU

handler.py

@@ -5,7 +5,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 class EndpointHandler():
     def __init__(self, path=""):
         torch.set_default_device("cuda")
-        self.model = AutoModelForCausalLM.from_pretrained("chuckfinca/arithmephi", torch_dtype="auto", trust_remote_code=True)
+        self.model = AutoModelForCausalLM.from_pretrained("chuckfinca/arithmephi", torch_dtype="auto", trust_remote_code=True, device_map = 'cuda')
         self.tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: