Udpate ordering

vllm-project · Jul 18, 2024 · 64382b3 · 64382b3
1 parent bc616e6
commit 64382b3
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/vllm/utils.py b/vllm/utils.py
@@ -663,10 +663,11 @@ def make_tensor_with_pad(
     np_dtype = TORCH_DTYPE_TO_NUMPY_DTYPE[dtype]
     padded_x = make_array_with_pad(x, max_len, pad, np_dtype)
 
-    tensor = torch.from_numpy(padded_x)
+    tensor = torch.from_numpy(padded_x).to(device)
     if pin_memory:
         tensor = tensor.pin_memory()
-    return tensor.to(device)
+
+    return tensor
 
 
 def async_tensor_h2d(