fix bugs in pytest

hpcaitech · Apr 3, 2024 · 93db9fc · 93db9fc
1 parent a801b9e
commit 93db9fc
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/colossalai/inference/kv_cache/kvcache_manager.py b/colossalai/inference/kv_cache/kvcache_manager.py
@@ -412,12 +412,12 @@ def free_block_tables(self, block_tables: torch.Tensor, first_n: int = None) ->
         """
         assert block_tables.dim() == 2
         first_n = block_tables.size(0) if first_n is None else first_n
-        for block_table in block_tables[:first_n].tolist():
+        for block_table in block_tables[:first_n]:
             self.free_block_table(block_table)
 
     def clear_all(self) -> None:
         """Clear all the references and allocations on all the cache blocks."""
-        for block in self._cache_blocks.tolist():
+        for block in self._cache_blocks:
             block.clear()
         self._available_blocks = self.num_blocks
         self._block_states[:] = 1

diff --git a/tests/test_infer/test_ops/cuda/test_silu_and_mul.py b/tests/test_infer/test_ops/cuda/test_silu_and_mul.py
@@ -20,7 +20,8 @@ def test_silu_and_mul(SHAPE_X, SHAPE_Y, SHAPE_Z, dtype):
     act_out = torch.nn.functional.silu(ref_input[0], inplace=True)
     ref_out = act_out * ref_input[1]
 
-    origin_out = inference_ops.silu_and_mul(origin_input)
+    origin_out = torch.empty_like(ref_out)
+    inference_ops.silu_and_mul(origin_input, origin_out)
 
     if dtype == torch.float32:
         assert torch.allclose(origin_out, ref_out, atol=1e-5, rtol=1e-5)