I am trying to implement a matrix factorization algorithm in pytorch. Specifically, I have a class with matrices A,B,C and want to train A and B such that AB = C. My first try was to write the training as
for i in range(self.max_iter):
            self.optimizer.zero_grad()
            loss = ((torch.mm(self.A, self.B) - self.C)**2).mean()
            loss.backward()
            self.optimizer.step()

The variables are kept as
class nmf_solver:

    def __init__(self, A,B,C, step_size=STEP_SIZE):

        """
        solves | AB - C |^2
        """

        self.A = Variable(A.cuda(), requires_grad=True)
        self.B = Variable(B.cuda(), requires_grad=True)
        self.C = C.cuda()

        self.step_size = step_size

        self.optimizer = optim.SGD([self.A,self.B], lr=self.step_size, momentum=0.9)

But I’m met with a “RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed.”. Since there isn’t any input to the matrix factorization algorithm, I don’t really know if nn.module is appropriate to subclass.
Is there any way to make this work in pytorch?