diff --git a/api/onnx_web/worker/command.py b/api/onnx_web/worker/command.py
index 089be538..1b0a1a0a 100644
--- a/api/onnx_web/worker/command.py
+++ b/api/onnx_web/worker/command.py
@@ -41,6 +41,7 @@ class JobCommand:
         args: Any,
         kwargs: dict[str, Any],
     ):
+        self.device = device
         self.name = name
         self.fn = fn
         self.args = args