Update handler.py

handler.py

@@ -2,15 +2,15 @@ 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
+# get dtype
+dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
 
 
 class EndpointHandler:
     def __init__(self, path=""):
         # load the model
         tokenizer = AutoTokenizer.from_pretrained(path)
-        model = AutoModelForCausalLM.from_pretrained(path, low_cpu_mem_usage=True)
+        model = AutoModelForCausalLM.from_pretrained(path, device_map="auto",torch_dtype=dtype)
         # create inference pipeline
         self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)