Update handler.py (#1)

- Update handler.py (e52e5bbe1f709006910d411d04e2b84270ce3224)


Co-authored-by: Philipp Schmid <philschmid@users.noreply.huggingface.co>

handler.py

@@ -1,7 +1,7 @@
 import torch
 from typing import Dict, List, Any
 # from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 # check for GPU
 device = 0 if torch.cuda.is_available() else -1
@@ -13,7 +13,7 @@ class EndpointHandler:
         tokenizer = AutoTokenizer.from_pretrained(path)
         # model = AutoModel.from_pretrained(path, low_cpu_mem_usage=True)
         # model = AutoModelForCausalLM.from_pretrained(path, low_cpu_mem_usage=True)
-        model = AutoModelForSeq2SeqLM.from_pretrained(path, low_cpu_mem_usage=True)
+        model = AutoModelForCausalLM.from_pretrained(path, low_cpu_mem_usage=True)
         # create inference pipeline
         self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)