Update handler.py

handler.py

@@ -161,8 +161,11 @@ class EndpointHandler:
         # load the model
         tokenizer = transformers.AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
         model = GPTJForCausalLM.from_pretrained(path, low_cpu_mem_usage=True)
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        
+        # check for GPU
+        device = 0 if torch.cuda.is_available() else -1
         model.to(device)
+        
         # create inference pipeline
         self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)