Unverified Commit 81ce2a4b authored by Woosuk Kwon's avatar Woosuk Kwon Committed by GitHub
Browse files

[Minor] Fix type annotation in Mixtral (#2036)

parent 5dd80d37
......@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment