memory leak in pytorch autograd function when assigning and returning context data

Viewed 115

My memory usage steadily increases when I run the program below.
One can avoid that by either not requiring a grad on the data or even by removing the *2 in the forward method.
Note that there is also no leak when one simply returns x*2 explicitly (i.e. reevaluating).
Is this in any way expected? What exactly happens if I assign data to ctx?

tested with torch version 1.8.1 (cpu only)

import torch
print(torch.__version__)

class LeakyBoi(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.x = x*2
        return ctx.x

EPOCHS = 1000000

data = torch.ones((100,100), requires_grad=True) # set False to avoid leak
lb = LeakyBoi()

for i in range(EPOCHS):
    lb.apply(data)
0 Answers
Related