Fix memory leak on masked Tensor

albanD · albanD · commit 41881d776a7b · 2024-10-14T09:50:16.000-04:00
diff --git a/test/test_maskedtensor.py b/test/test_maskedtensor.py
@@ -78,10 +78,6 @@ def _compare_forward_backward(data, mask, fn):
     _compare_mt_t(masked_res, tensor_res)
     _compare_mt_t(mt.grad, t.grad, atol=1e-06)
 
-    # Free up the masked tensors manually to avoid memory leak
-    del mt._masked_mask
-    del mt._masked_data
-
 def _create_random_mask(shape, device):
     return make_tensor(shape, device=device, dtype=torch.bool)
 
@@ -229,11 +225,6 @@ def test_stack(self, device):
         for mt, t in zip(masked_tensors, data_tensors):
             _compare_mt_t(mt.grad, t.grad, atol=1e-06)
 
-        # Free up the masked tensors manually to avoid memory leak
-        for mt in masked_tensors:
-            del mt._masked_mask
-            del mt._masked_data
-
     def test_to_sparse(self, device):
         for sample in _generate_sample_data(device=device):
             data = sample.input
diff --git a/torch/masked/maskedtensor/core.py b/torch/masked/maskedtensor/core.py
@@ -334,7 +334,7 @@ def get_data(self):
         class GetData(torch.autograd.Function):
             @staticmethod
             def forward(ctx, self):
-                return self._masked_data
+                return self._masked_data.detach()
 
             @staticmethod
             def backward(ctx, grad_output):
