[TPU] Skip creating empty tensor (vllm-project#7630)

fixie-ai · Aug 17, 2024 · ce14335 · ce14335
1 parent bbf55c4
commit ce14335
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/vllm/worker/tpu_worker.py b/vllm/worker/tpu_worker.py
@@ -271,7 +271,10 @@ def _make_src_to_dst(
     mapping: List[Tuple[int, int]],
     src_device: Union[torch.device, str],
     dst_device: Union[torch.device, str],
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
+    if not mapping:
+        return None
+
     src_indices = [i for i, _ in mapping]
     dst_indices = [i for _, i in mapping]
     src_indices = torch.tensor(src_indices,