Update handler.py

handler.py

@@ -43,22 +43,21 @@ class EndpointHandler:
         )
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
-        """
-        data 格式：
-          {
-            "inputs": "your prompt here"
-          }
-        """
         prompt = data["inputs"]
-
-        # ➡️ 只把输入张量放到 cuda:0（与模型第一层同卡）
-        inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0")
-
-        # 生成
+    
+        # ① 自动抓 embedding 所在 GPU
+        first_device = next(self.model.parameters()).device
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(first_device)
+    
+        # ② 生成（其余逻辑不变）
         with torch.inference_mode():
-            output_ids = self.model.generate(**inputs, **self.generation_kwargs)
-
-        generated_text = self.tokenizer.decode(
-            output_ids[0], skip_special_tokens=True
-        )
-        return {"generated_text": generated_text}
+            output_ids = self.model.generate(
+                **inputs,
+                **self.generation_kwargs,
+            )
+    
+        return {
+            "generated_text": self.tokenizer.decode(
+                output_ids[0], skip_special_tokens=True
+            )
+        }