Reyad-Ahmmed
/

getvars-generic

Model card Files Files and versions

Reyad-Ahmmed commited on Feb 9, 2025

Commit

808c660

·

verified ·

1 Parent(s): 262130c

Update handler.py

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -25,7 +25,7 @@ class EndpointHandler:
         model_name = model_dir2 #"./json_extraction_all"  # Pretrained model for sentiment analysis
         self.tokenizer = T5Tokenizer.from_pretrained(model_name)
-        self.model = T5ForConditionalGeneration.from_pretrained(model_name)
         self.model.eval()  # Set model to evaluation mode (no training)
         # Check if the model is on GPU

         model_name = model_dir2 #"./json_extraction_all"  # Pretrained model for sentiment analysis
         self.tokenizer = T5Tokenizer.from_pretrained(model_name)
+        self.model = T5ForConditionalGeneration.from_pretrained(model_name, device_map="auto")
         self.model.eval()  # Set model to evaluation mode (no training)
         # Check if the model is on GPU