Repository: pytorch
Issue: Cuda memory leak check is somehow unstable with repeat_test_for_types
Labels: module: memory usage,triaged,module: testing

## 🐛 Bug

Cuda memory leak check is somehow unstable with repeat_test_for_types

## To Reproduce

Name the following file to `test_AAA.py`, put it into the `test/` folder, simply run it with `IN_CI= python test/test_AAA.py -v`


```python
import torch
from torch.testing._internal.common_utils import run_tests, TestCase, repeat_test_for_types

class TestAAA(TestCase):
    _do_cuda_memory_leak_check = True

    @repeat_test_for_types([torch.float])
    def test_memory_leak_cuda(self, dtype=torch.float):
        one = torch.ones(4, dtype=dtype, device='cuda')
        two = torch.ones(4, dtype=dtype, device='cuda') * 2
        self.assertEqual(one, two)

    def test_memory_leak_cuda__no_repeat(self, dtype=torch.float):
        one = torch.ones(4, dtype=dtype, device='cuda')
        two = torch.ones(4, dtype=dtype, device='cuda') * 2
        self.assertEqual(one, two)

if __name__ == '__main__':
    run_tests()
```

Sample error message

```
$ IN_CI= python test/test_AAA.py -v
test_memory_leak_cuda (__main__.TestAAA) ... FAIL
test_memory_leak_cuda__no_repeat (__main__.TestAAA) ... FAIL

======================================================================
FAIL: test_memory_leak_cuda (__main__.TestAAA) (dtype=torch.float32)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/xwang/Developer/pytorch/torch/testing/_internal/common_utils.py", line 230, in call_helper
    f(self, *args, dtype=dtype)
  File "/home/xwang/Developer/pytorch/test/test_AAA.py", line 11, in test_memory_leak_cuda
    self.assertEqual(one, two)
  File "/home/xwang/Developer/pytorch/torch/testing/_internal/common_utils.py", line 1275, in assertEqual
    super().assertTrue(result, msg=self._get_assert_msg(msg, debug_msg=debug_msg))
AssertionError: False is not true : Tensors failed to compare as equal!With rtol=1.3e-06 and atol=1e-05, found 4 element(s) (out of 4) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 1.0 (1.0 vs. 2.0), which occurred at index 0.

======================================================================
FAIL: test_memory_leak_cuda (__main__.TestAAA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/xwang/Developer/pytorch/torch/testing/_internal/common_utils.py", line 979, in wrapper
    method(*args, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/testing/_internal/common_utils.py", line 767, in __exit__
    self.testcase.assertEqual(
  File "/home/xwang/Developer/pytorch/torch/testing/_internal/common_utils.py", line 1320, in assertEqual
    super().assertTrue(result, msg=self._get_assert_msg(msg, debug_msg=debug_msg))
AssertionError: False is not true : Scalars failed to compare as equal! Comparing 0 and 1024 gives a difference of 1024, but the allowed difference with rtol=0 and atol=0 is only 0!
__main__.TestAAA.test_memory_leak_cuda leaked 1024 bytes CUDA memory on device 0

======================================================================
FAIL: test_memory_leak_cuda__no_repeat (__main__.TestAAA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/xwang/Developer/pytorch/torch/testing/_internal/common_utils.py", line 979, in wrapper
    method(*args, **kwargs)
  File "/home/xwang/Developer/pytorch/test/test_AAA.py", line 16, in test_memory_leak_cuda__no_repeat
    self.assertEqual(one, two)
  File "/home/xwang/Developer/pytorch/torch/testing/_internal/common_utils.py", line 1275, in assertEqual
    super().assertTrue(result, msg=self._get_assert_msg(msg, debug_msg=debug_msg))
AssertionError: False is not true : Tensors failed to compare as equal!With rtol=1.3e-06 and atol=1e-05, found 4 element(s) (out of 4) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 1.0 (1.0 vs. 2.0), which occurred at index 0.

----------------------------------------------------------------------
Ran 2 tests in 1.477s

FAILED (failures=3)
```

Note that: 

```
__main__.TestAAA.test_memory_leak_cuda leaked 1024 bytes CUDA memory on device 0
``` 

indicates a cuda memory leak. However, I think this is a very false positive (?) thing? `self.assertEqual` reports the mismatch and throws python exception to terminate the program. That makes the python tensor not GC'ed or the C++ tensor not destroyed. Finally, the `CudaMemoryLeakCheck` context manager checks the memory usage and reports this error.

https://github.com/pytorch/pytorch/blob/6d87b3667f0f799939f74b88aca49ddf7d4bfe5a/torch/testing/_internal/common_utils.py#L759-L769

This "memory leak" report only exists when `repeat_test_for_types` is used. Is it because this decorator hides the exception info from `CudaMemoryLeakCheck`?

https://github.com/pytorch/pytorch/blob/6d87b3667f0f799939f74b88aca49ddf7d4bfe5a/torch/testing/_internal/common_utils.py#L224-L233

## Expected behavior

There are many tests wrapped with `repeat_test_for_types` in pytorch. Test failures with memory leak could be thrown in these tests and make users confused.

## Environment

```
Collecting environment information...
PyTorch version: 1.9.0a0+git02af4b5
Is debug build: False
CUDA used to build PyTorch: 11.2
ROCM used to build PyTorch: N/A

OS: Manjaro Linux (x86_64)
GCC version: (GCC) 10.2.0
Clang version: Could not collect
CMake version: version 3.19.3

Python version: 3.9 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 11.2.67
GPU models and configuration: 
GPU 0: GeForce RTX 2070 SUPER
GPU 1: GeForce GTX 1070 Ti

Nvidia driver version: 460.32.03
cuDNN version: Probably one of the following:
/usr/lib/libcudnn.so.8.0.5
/usr/lib/libcudnn_adv_infer.so.8.0.5
/usr/lib/libcudnn_adv_train.so.8.0.5
/usr/lib/libcudnn_cnn_infer.so.8.0.5
/usr/lib/libcudnn_cnn_train.so.8.0.5
/usr/lib/libcudnn_ops_infer.so.8.0.5
/usr/lib/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.9.0a0+git02af4b5
[pip3] torchvision==0.10.0a0+7d955df
[conda] Could not collect
```

## Additional context

N/A

cc @ptrblck @ngimel @malfet