alvarobartt
/

Magma-8B

Image-Text-to-Text

text-generation

Model card Files Files and versions

alvarobartt HF Staff commited on May 8

Commit

f74a199

·

verified ·

1 Parent(s): cb5f6e9

Update handler.py

Files changed (1) hide show

handler.py +1 -0

handler.py CHANGED Viewed

@@ -18,6 +18,7 @@ class EndpointHandler:
         self.model = AutoModelForCausalLM.from_pretrained(
             model_dir, trust_remote_code=True, torch_dtype=torch.bfloat16
         ).eval()
         self.processor = AutoProcessor.from_pretrained(
             model_dir, trust_remote_code=True
         )

         self.model = AutoModelForCausalLM.from_pretrained(
             model_dir, trust_remote_code=True, torch_dtype=torch.bfloat16
         ).eval()
+        self.model.to("cuda")
         self.processor = AutoProcessor.from_pretrained(
             model_dir, trust_remote_code=True
         )