Memory (CPU and GPU) leaks during the 1st epoch
🐛 Bug
Hello.
This memory leak occurs during the first epoch. If one has a large epoch time (I had > 10 days), the OOM error will come. It's interesting, that in precision=16 mode, it leaks out on the GPU and the CPU both. If we switch amp optimization off (precision=32), the leak goes only on the CPU.
Also, I checked the number of tensors, which are tracked by the garbage collector. And it appeared to be linearly increasing during the first epoch, and then (on the 2nd epoch starts), it falls to the initial value and begins increasing again.
Let me provide the plots:

Experiment 1: amp_level='O2', precision=16

The number of tensors, tracked by garbage collector

GPU (the 2nd in my case) usage, tracked by pytorch-lightning

CPU memory usage by the process (bytes)

Experiment 2: amp_level=None, precision=None

The number of tensors, tracked by garbage collector

GPU (the 2nd in my case) usage, tracked by pytorch-lightning

CPU memory usage by the process (bytes)

As you can see, both cases have a CPU leak. The "amp"-case also has a GPU leak.
Also, it's clear, that such leaky behavior stops when the 2nd epoch starts.
On these plots, the 2nd epoch starts on the 2nd "saw claw" of the "Num-of-tensors" plot.
Also, there is another observation: the speed of tensors number increasing is 1001. And this is my forward pass method:
    def training_step(self, batch, batch_idx):
        losses = self.forward(batch)
        num_of_tensors = get_num_of_tensors()
        log = {'Num-of-tensors': num_of_tensors, 'Cpu-mem-usg': get_cpu_mem()}

        for i, loss in enumerate(losses):
            log[f'loss{i}'] = loss

        print(num_of_tensors)
        return {'loss': losses[0], 'log': log}
Here I return exactly 1001 tensor: one for loss and 1000 for log.
In my real experiments I had only 3 tensors. It took ~2-3 days to get OOM. But in the current example (see To Reproduce) it will crash much faster.
To Reproduce
Steps to reproduce the behavior:

Execute Code sample (this script has no arguments, so change needed values manually in script).
Go to the tensorboard to check plots.

Code sample
https://gist.github.com/alexeykarnachev/47de06b93a717ab0664eded42ed2826a
Expected behavior
The number of tensors, GPU and CPU memory does not increase during the training.
Environment
PyTorch version: 1.4.0
OS: Ubuntu 16.04.6 LTS
Python version: 3.7
Versions of relevant libraries:
[pip] numpy==1.18.1
[pip] pytorch-lightning==0.7.3
[pip] torch==1.4.0
[pip] torchvision==0.5.0
Additional context
Sorry for so messy flow of the information, but I don't know, how to structure it more clearly.