From 64b77dfd7e1378853ec7b189f3d7d0e51ce18855 Mon Sep 17 00:00:00 2001 From: Kunshang Ji Date: Fri, 10 May 2024 20:52:48 +0800 Subject: [PATCH] [Core]fix type annotation for `swap_blocks` (#4726) --- vllm/_custom_ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 35a9f6329fc42..42dedfdf76c4f 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional, Tuple +from typing import Optional, Tuple import torch @@ -266,7 +266,7 @@ def copy_blocks(key_caches: torch.Tensor, value_caches: torch.Tensor, def swap_blocks(src: torch.Tensor, dst: torch.Tensor, - block_mapping: Dict[int, int]) -> None: + block_mapping: torch.Tensor) -> None: vllm_cache_ops.swap_blocks(src, dst, block_mapping)