JuanFMontesinos:

My guess is that defining the constant outside nn.Module breaks the graph. Once you define them inside nn.Module they are gonna be registered as buffers and it sounds bad to me.
Why don’t you just define them as regular nn.Parameters and pass the nn.Module to the optimizer?
Anyway, can you provide a runnable script?
WRT the inplace ops, indeed that M[…,2] is not allowed.


Thank you.
Please note that I manage to learn rotation angle while the parameters are passed to the nn.Module so this doesn’t seem to be the problem.
Attached is a runbable code (originally a notebook):
import kornia; print(kornia.__version__)


# In[2]:


import torch
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

bs= 4#096
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=bs, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=bs, shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')


# In[3]:


from torch import nn
class DRotation(nn.Module):
    def __init__(self, angle):
        super(DRotation, self).__init__()
        self.angle = angle
        
        
    def forward(self, input):
        if self.angle.shape[0] != input.shape[0]:
            angle = self.angle.repeat(input.shape[0])
        else:
            angle = self.angle
        
        # 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)

        # apply the transformation to original image
        _, _, h, w = input.shape
        out = kornia.warp_affine(input, M, dsize=(h, w))
        
        return out


# In[4]:


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
#         shape = list(M.shape)
#         shape[-1] -= 1
#         M = M + torch.cat([torch.zeros(shape), translations.unsqueeze(-1)], 2)

        
        # apply the transformation to original image
        out = kornia.warp_affine(input, M, dsize=(h, w))
        
        return out


# In[5]:


from torch.nn.parameter import Parameter

class LearnableAug(nn.Module):
    def __init__(self, aug_names=[]):
        super(LearnableAug, self).__init__()
        
        self.fixed_transforms = []
        self.learnable_transforms = []
        self.have_entropy = []
        self._param_groups = []
        self.lrs = []
        self.bounds= {}
        self.eps = 0 #1e-6
        
        if 'translation' in aug_names:
            var = torch.tensor([0.3], dtype=torch.float32, requires_grad=False)
            self.register_buffer('_x_translation', var)
            var = torch.tensor([-0.2], dtype=torch.float32, requires_grad=False)
            self.register_buffer('_y_translation', var)
            
            self.fixed_transforms.append(DTranslation(
                x_translation=self._buffers['_x_translation'],
                y_translation=self._buffers['_y_translation'],
            ))

            var = torch.tensor([0], dtype=torch.float32)
            param = Parameter(var, requires_grad=True)
            self.register_parameter(name='x_translation', param=param)
            
            var = torch.tensor([0], dtype=torch.float32)
            param = Parameter(var, requires_grad=True)
            self.register_parameter(name='y_translation', param=param)
            
            self.learnable_transforms.append(DTranslation(
                x_translation=self._parameters['x_translation'],
                y_translation=self._parameters['y_translation'],
            ))
            
            self._param_groups.append(['x_translation'])
            self.lrs.append(1)
            self.bounds['x_translation'] = (-0.5, 0.5)
            
            self._param_groups.append(['y_translation'])
            self.lrs.append(1)
            self.bounds['y_translation'] = (-0.5, 0.5)
        
            
        if 'rotation' in aug_names and 'fixed_affine' not in aug_names:
            # Fixed Rotation
            var = torch.tensor([15], dtype=torch.float32, requires_grad=False)
            self.register_buffer('_angle', var)
            self.fixed_transforms.append(DRotation(var))

            var = torch.tensor([0], dtype=torch.float32)
            param = Parameter(var, requires_grad=True)
            self.register_parameter(name='angle', param=param)
            self.learnable_transforms.append(DRotation(param))
            self._param_groups.append(['angle'])
            self.lrs.append(1)
            self.bounds['angle'] = (-60, 60)
       
            
        self.fixed_transforms = torch.nn.Sequential(*self.fixed_transforms)
        self.learnable_transforms = torch.nn.Sequential(*self.learnable_transforms)

    def apply_constraints(self):
        for param_name, bounds in self.bounds.items():
            if self._parameters[param_name] < bounds[0]:
                if isinstance(bounds[0], torch.Tensor):
                    self._parameters[param_name].data = bounds[0].data
                else:
                    bound_tensor = torch.ones_like(self._parameters[param_name]) * bounds[0] 
                    self._parameters[param_name].data = bound_tensor.data
                    
            if self._parameters[param_name] > bounds[1]:
                if isinstance(bounds[1], torch.Tensor):
                    self._parameters[param_name].data = bounds[1].data
                else:
                    bound_tensor = torch.ones_like(self._parameters[param_name]) * bounds[1] 
                    self._parameters[param_name].data = bound_tensor.data
                    
                    
    def forward(self, x, fixed=False):
        if not fixed:
            self.apply_constraints()
            out = self.learnable_transforms(x)
        else:
            out = self.fixed_transforms(x)
        return out
    
    @property
    def param_groups(self):
        return self._param_groups
    
    def get_params_for_optimizer(self):
        param_groups = []
        for param_group, lr in zip(self._param_groups, self.lrs):
            param_groups.append(dict(params=[self._parameters[name] for name in param_group], lr=lr))
        
        return param_groups
    


# In[8]:


from IPython.display import clear_output

from time import sleep
from torch import optim
device = torch.device('cuda')

dataloader = trainloader

augs = LearnableAug(['rotation'])
# augs = LearnableAug(['translation'])
criterion = nn.MSELoss()
optimizer = optim.Adam(augs.get_params_for_optimizer())
scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10)#, T_mult)

iters = len(dataloader)
epochs = 1

for epoch in range(epochs):
    for i, data in enumerate(dataloader):
        images, labels = data

        # perform the transforms
        ft_images = augs(images, fixed=True)
        dt_images = augs(images)

        optimizer.zero_grad()
        loss = criterion(ft_images, dt_images)
        print(loss)
        loss.backward()
            
        optimizer.step()
        scheduler.step((epoch * iters + i) / (epochs * iters))