Update handler.py

handler.py

@@ -11,7 +11,19 @@ class EndpointHandler:
             model, checkpoint=model_dir, device_map="auto"
         )                      # 自动跨 GPU 切层
     def __call__(self, data):
-        prompt = data["inputs"]
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
-        out_ids = self.model.generate(**inputs, max_new_tokens=256)
-        return {"generated_text": self.tokenizer.decode(out_ids[0], skip_special_tokens=True)}
+    prompt = data["inputs"]
+
+    inputs = self.tokenizer(
+        prompt, return_tensors="pt"
+    ).to("cuda:0")          # 👈 把 input_ids/attention_mask 都放到 0 号卡
+
+    out_ids = self.model.generate(
+        **inputs,
+        max_new_tokens=256,
+    )
+    return {
+        "generated_text": self.tokenizer.decode(
+            out_ids[0], skip_special_tokens=True
+        )
+    }
+