Remove force data release for CPU offloading.

Signed-off-by: Dennis Liu <[email protected]>
NVIDIA · Feb 13, 2025 · 594ea31 · 594ea31
1 parent 3ffd732
commit 594ea31
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py
@@ -381,7 +381,6 @@ def bulk_offload_group(self, group_to_offload):
                     if self.tensor_need_offloading_checker(tensor_on_device):
                         state = SynchronizedGroupOffloadHandler.offload(tensor_on_device)
                         self.tensor_tag_to_state[tensor_tag] = state
-                        tensor_on_device.data = torch.Tensor()  # Force to release memory
 
     def synchronize_on_group_commit_forward(self, current_group):
         """Synchronize on group commit forward."""