GPU memory leak in For Loop with AMP mode
🐛 Bug
Computation in For loop with AMP ON cause GPU memory leak, and crash training.
To Reproduce
Notebook.
Core part:
    def generate(self, device: str) -> None:
        ipt = torch.tensor([[float(j) for j in range(0, self.size_i)]], device=device)
        for _ in range(0, 10000):
            self.fc(ipt)
Run .generate() in validation step with precision=16 trainer's argument.
Expected behavior
self.fc(ipt)'s output is simply discarded, so GPU memory will be released.
GPU memory leak never occur.
Actual behavior
Out of Memory error only when precision=16, not precision=32.
RuntimeError: CUDA out of memory. ....

Environment
Check in Colab Notebook.