I’m trying to accumulate into a variable using a loop, as follows. In each iteration of the loop, I do some computation using the variables, and add the result to total_act. Note that the actual operation I’m trying to do is more complicated; this is just a minimalistic example to reproduce the problem.
import torch
import torch.nn as nn
from torch.nn import Parameter

num_x = 20000
num_y = 1000
emb_dim = 500

class Model(torch.nn.Module):

    def __init__(self):
        super(Model, self).__init__()
        self.x_embed = Parameter(torch.FloatTensor(num_x, emb_dim))
        self.y_embed = Parameter(torch.FloatTensor(num_y, emb_dim))

        self.w = Parameter(torch.FloatTensor(emb_dim, emb_dim))

    def forward(self):
        total_act = 0

        for i in range(num_y):
            import pdb; pdb.set_trace()

            trans = self.x_embed - self.y_embed[i]
            trans = torch.mm(trans, self.w)  # If we remove this line, the problem doesn't occur

            total_act += trans

        return total_act

model = Model()
model = model.cuda()

final_act = model.forward()

Now what I notice is that with each iteration of the for loop, the memory occupied on the GPU keeps increasing. I’ve put a pdb in the loop so that one can track memory consumption for each iteration.
My guess is that this is due to the intermediary states of the various variables occupying space, as though it’s building one graph for every single iteration of the loop. How do I solve this problem?