Tokenizer to xpu

mudler · Mar 7, 2024 · 6703787 · 6703787
1 parent 9ff1cf8
commit 6703787
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/backend/python/transformers/transformers_server.py b/backend/python/transformers/transformers_server.py
@@ -160,6 +160,8 @@ def Predict(self, request, context):
         inputs = self.tokenizer(request.Prompt, return_tensors="pt").input_ids
         if self.CUDA:
             inputs = inputs.to("cuda")
+        if XPU:
+            inputs = inputs.to("xpu")
 
         outputs = self.model.generate(inputs,max_new_tokens=max_tokens, temperature=request.Temperature, top_p=request.TopP)