From 81ce2a4b26c7a76742039d7b96882e2e741ac5c8 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 11 Dec 2023 11:32:39 -0800 Subject: [PATCH] [Minor] Fix type annotation in Mixtral (#2036) --- vllm/model_executor/models/mixtral.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/models/mixtral.py b/vllm/model_executor/models/mixtral.py index 37afccda7f01..f14821547270 100644 --- a/vllm/model_executor/models/mixtral.py +++ b/vllm/model_executor/models/mixtral.py @@ -251,7 +251,7 @@ class BlockSparseMoE(nn.Module): return column_indices_t, offsets_t, block_offsets_t def topology(self, x: torch.Tensor, - padded_bins: torch.Tensor) -> stk.Matrix: + padded_bins: torch.Tensor) -> "stk.Matrix": padded_tokens, _ = x.size() assert padded_tokens % self.blocking == 0 assert self.ffn_dim_per_partition % self.blocking == 0