alexadam
/

handler_trial

Model card Files Files and versions

alexadam commited on May 30, 2024

Commit

8a6e1d9

·

verified ·

1 Parent(s): 826348d

Update handler.py

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -7,7 +7,7 @@ import torch
 class EndpointHandler():
     def __init__(self, path=""):
         # load the model
-        self.model = transformers.AutoModelForCausalLM.from_pretrained(
             "gpt2", torch_dtype=torch.float16, output_hidden_states=True
         )
         self.model = self.model.cuda()

 class EndpointHandler():
     def __init__(self, path=""):
         # load the model
+        self.model = AutoModelForCausalLM.from_pretrained(
             "gpt2", torch_dtype=torch.float16, output_hidden_states=True
         )
         self.model = self.model.cuda()