Fix test_deallocate_output_tensor

microsoft · Jun 13, 2024 · f31c94f · f31c94f
1 parent 7eb36a1
commit f31c94f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/pipeline_parallel/test_schedules.py b/tests/pipeline_parallel/test_schedules.py
@@ -21,7 +21,9 @@ def test_get_forward_backward_func():
 def test_deallocate_output_tensor():
     out = torch.tensor([[1, 2, 3], [4, 5, 6]])
     schedule.deallocate_output_tensor(out)
-    assert(out.nelement() == 1) 
+    assert(out.nelement() == 6)
+    schedule.deallocate_output_tensor(out, True)
+    assert(out.nelement() == 1)
 
 def test_forward_backward_func_without_pipeline_parallel(mocker):
     from megatron.core.pipeline_parallel import get_forward_backward_func