You are recreating the theta tensor in the forward, which will detach the parameters.
You could either create theta in the __init__ method as:
        self.theta = nn.Parameter(torch.tensor([
            [self.scalew, 0, self.transX],
            [0, self.scaleh, self.transY]
        ])[None])

or use torch.cat or torch.stack to create theta in the forward method from the parameters.