Hi,
I have been trying to learn translation (x, y) parameters with Kornia in the following manner:
class DTranslation(nn.Module):
    def __init__(self, x_translation, y_translation):
        super(DTranslation, self).__init__()
        self.translations = torch.stack([x_translation, y_translation], 1)
        self.angle = torch.tensor([0])
        
        
    def forward(self, input):
        _, _, h, w = input.shape
        if self.angle.shape[0] != input.shape[0]:
            angle = self.angle.repeat(input.shape[0])
        else:
            angle = self.angle
            
        if self.translations.shape[0] != input.shape[0]:
            translations = self.translations.repeat([input.shape[0], 1])
        else:
            translations = self.translations
        
        translations[:, 0] *= h
        translations[:, 1] *= w
        
        # define the rotation center
        center = torch.ones(2)
        center[..., 0] = input.shape[3] / 2  # x
        center[..., 1] = input.shape[2] / 2  # y
        center = center.repeat(input.shape[0], 1)

        # define the scale factor
        scale = torch.ones(input.shape[0])

        # compute the transformation matrix
        M = kornia.get_rotation_matrix2d(center, -angle, scale)
        
        # Translate
        M[..., 2] += translations  # tx/ty
        
        # apply the transformation to original image
        out = kornia.warp_affine(input, M, dsize=(h, w))
        
        return out

tx = torch.tensor([0.3], dtype=torch.float32)
tx_p = Parameter(tx, requires_grad=True)
            
ty = torch.tensor([-0.2], dtype=torch.float32)
ty_p = Parameter(ty, requires_grad=True)

translation = DTranslation(x_translation=tx_p, y_translation=ty_p)

criterion = nn.MSELoss()
optimizer = optim.Adam([tx_p, ty_p], lr=1)

for x, y in dataloader: 
        optimizer.zero_grad()
        loss = criterion(x, translation(x))
        loss.backward()
        optimizer.step()

The first backward call passes, but the second one fails:
RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed. Specify retain_graph=True when calling backward the first time.

When I follow the instructions and instead I use:
 loss.backward(retain_graph=True)

I receive the following error:
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [1]] is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

I also tried to avoid the inplace operation:
        M[..., 2] += translations  # tx/ty

and use instead:
shape = list(M.shape)
shape[-1] -= 1
M = M + torch.cat([torch.zeros(shape), translations.unsqueeze(-1)], 2)

but got the same errors.
Maybe another inplace operation causes this ? e.g. https://github.com/kornia/kornia/blob/5a736409a9a133da27c3dfa581bba2bc71f27286/kornia/geometry/conversions.py#L122 4
Or is it something else ?
It is worth mentioning that I do manage to backpropagated through rotation angle and shear (x, y) but only translation seems to be the problem, i.e. when I comment out
        M[..., 2] += translations  # tx/ty

No errors occur, but of course I cannot learn the translation parameters either.
Any thoughts ?