DW-ReCo
/

spot_llama3_training_ds_v15.2_lora

text-generation-inference

Inference Endpoints

Model card Files Files and versions Community

DW-ReCo commited on Aug 26

Commit

ff3c462

•

1 Parent(s): 80ec456

Update handler.py

Files changed (1) hide show

handler.py +2 -0

handler.py CHANGED Viewed

@@ -36,6 +36,8 @@ class EndpointHandler:
                     "",  # output - leave this blank for generation!
                 )
             ], return_tensors="pt")
         outputs = self.model.generate(**inputs,
                     max_new_tokens=1048,

                     "",  # output - leave this blank for generation!
                 )
             ], return_tensors="pt")
+        inputs = inputs.to('cuda')
         outputs = self.model.generate(**inputs,
                     max_new_tokens=1048,