[DeepSeekV32] Bug fix to ensure page_table and result in same type (#12300)
This commit is contained in:
@@ -103,7 +103,7 @@ def transform_index_page_table_decode_ref(
|
|||||||
result = torch.empty_like(topk_indices, dtype=torch.int32)
|
result = torch.empty_like(topk_indices, dtype=torch.int32)
|
||||||
assert result.shape == topk_indices.shape
|
assert result.shape == topk_indices.shape
|
||||||
torch.gather(
|
torch.gather(
|
||||||
page_table,
|
page_table.to(result.dtype),
|
||||||
dim=1,
|
dim=1,
|
||||||
index=topk_indices.clamp(min=0),
|
index=topk_indices.clamp(min=0),
|
||||||
out=result,
|
out=result,
|
||||||
|
|||||||
Reference in New Issue
Block a user