fredaddy
/

MiniCPM-v-2_6

Image-Text-to-Text

feature-extraction

Model card Files Files and versions

fredaddy commited on Sep 19, 2024

Commit

a1c2e19

·

verified ·

1 Parent(s): db6f154

Update handler.py

Files changed (1) hide show

handler.py +1 -2

handler.py CHANGED Viewed

@@ -14,8 +14,7 @@ class EndpointHandler:
             trust_remote_code=True,
             attn_implementation='sdpa',
             torch_dtype=torch.bfloat16 if self.device.type == "cuda" else torch.float32,
-            device_map="auto",
-        )
         self.model.eval()
         # Load the tokenizer

             trust_remote_code=True,
             attn_implementation='sdpa',
             torch_dtype=torch.bfloat16 if self.device.type == "cuda" else torch.float32,
+        ).to(self.device)
         self.model.eval()
         # Load the tokenizer