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)