ahhhh forgot an important conditional

pytorch · Nov 1, 2024 · a326ab7 · a326ab7
1 parent 6f1d03f
commit a326ab7
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/torchtune/training/_activation_offloading.py b/torchtune/training/_activation_offloading.py
@@ -314,8 +314,11 @@ def hook(outputs, inputs):
                         # non-deterministic (thus higher) memory usage, but this case
                         # should not happen often.
                         unpacked_tensor = self.bwd_tensor_stash[unpack_tensor_id]
-                        if torch._C._storage_Use_Count(
-                            unpacked_tensor.untyped_storage()._cdata
+                        if (
+                            torch._C._storage_Use_Count(
+                                unpacked_tensor.untyped_storage()._cdata
+                            )
+                            > storage_refcount
                         ):
                             unpacked_tensor.record_stream(self.s0)
                             del self.bwd_tensor_stash[unpack_tensor_id]