Sure! Here’s the full code snippet for computing the loss (a fully executable snippet might be too long, but I can provide it if needed). I’m generating a sequence of images by applying the same transformation n times and evaluating the reconstruction loss on each frame. I want the model to fix the latent representation of the first frame and reconstruct the successive frame by twiddling the fixed latent rep using trans_par:
    for epoch in range(epoch):
        for batch_idx, (data, target) in enumerate(train_loader):
            loss = 0
            optimizer.zero_grad()
            for i in range(movie_len):
                if i == 0:
                    transform = rand.choice(transform_set)
                    prev_frame = curr_frame = data
                    curr_frame = curr_frame.flatten(1).to(device)
                    output, latent_rep, trans_par = model(curr_frame)
                    latent_rep = latent_rep.detach().clone()
                    latent_rep = latent_rep.to(device)
                else:
                    curr_frame = transform(prev_frame)
                    prev_frame = curr_frame
                    curr_frame = curr_frame.flatten(1).to(device)
                    output, out_rep, trans_par = model(curr_frame, latent_rep)
                loss += (beta**i)*(F.mse_loss(output, curr_frame) + 5e-2*(1/B)*torch.norm(trans_par,1))
            loss.backward()
            optimizer.step()


I’ve tested it without reusing latent_rep(model defaults to generating a latent rep each frame) and it works fine. Only when I try to reuse latent_rep do I get the error.
Another cause might be the forward pass of my transformation module, where I concatenate the two inputs. Ideally, since they’re two “constant” tensors then it shouldn’t matter, but maybe I’m missing something?
class Transnet(nn.Module):
    def __init__(self, og_dim, latent_dim, trans_dim, k_sparse):
        super(Transnet,self).__init__()
        assert latent_dim <= og_dim, 'latent space must have lower dimension'
        assert trans_dim <= latent_dim, 'translation dimension must be subspace'
        self.og_dim = og_dim
        self.latent_dim = latent_dim
        self.trans_dim = trans_dim
        ttl_dim = og_dim + latent_dim
        self.ttl_dim = ttl_dim
        self.k_sparse = k_sparse
        self.fc1 = nn.Linear(ttl_dim, max(latent_dim, ttl_dim//16))
        self.fc2 = nn.Linear(max(latent_dim, ttl_dim//16), max(latent_dim, ttl_dim//32))
        self.fc3 = nn.Linear(max(latent_dim, ttl_dim//32), trans_dim)
        
    def forward(self,x,x0):
        x1 = torch.cat((x,x0),dim = 1) #create (B, N+M) tensor
        x1 = self.fc1(x1)
        x1 = F.relu(x1)
        x1 = self.fc2(x1)
        x1 = F.relu(x1)
        x1 = self.fc3(x1)
        x0[:,:self.trans_dim] += x1
        return x0, x1