I have a model in which I create a rotational matrix as a function of angle and would like to have angle as a trainable parameter of the model. However, angle is not updated, although it has requires_grad=True.
class Rotation(Affine):
    def __init__(self,scale,center=None, angle=None):
        super().__init__(scale,center)
        if angle is None:
          self.angle=nn.Parameter(torch.rand(1))
        else: self.angle=nn.Parameter(torch.tensor(angle,dtype=torch.float))
        self.shift=torch.zeros((1,1,2,1))

    def forward(self,xy):
        def rotational_matrix(angle):
            return torch.tensor([[torch.cos(angle), -torch.sin(angle)], [torch.sin(angle), 
                       torch.cos(angle)]],requires_grad=True).view(1, 1, 2, 2)
        return self.local_transformation(xy,self.global_transformation(xy,rotational_matrix(self.angle),self.shift))

What is the proper way to make sure that angle is being updated?