Update handler.py

handler.py

@@ -1,7 +1,7 @@
 import torch
 
 from typing import Any, Dict
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 
 
 class EndpointHandler:
@@ -9,7 +9,7 @@ class EndpointHandler:
         # load model and tokenizer from path
         self.tokenizer = AutoTokenizer.from_pretrained(path) # AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
 
-        config = transformers.AutoConfig.from_pretrained(name, trust_remote_code=True)
+        config = AutoConfig.from_pretrained(name, trust_remote_code=True)
         config.attn_config['attn_impl'] = 'triton'
         config.init_device = 'cuda:0' # For fast initialization directly on GPU!
         config.max_seq_len = 4096 # (input + output) tokens can now be up to 4096