Update handler.py

handler.py

@@ -2,6 +2,10 @@ import torch
 from typing import Dict, List, Any
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
+# check for GPU
+device = 0 if torch.cuda.is_available() else -1
+
+
 format_input = (
     "Below is an instruction that describes a task. "
     "Write a response that appropriately completes the request.\n\n"
@@ -23,6 +27,7 @@ class EndpointHandler:
             "text-generation",
             model=model,
             tokenizer=tokenizer,
+            device=device,
             max_length=256,
         )