In order to prevent the transformations shearing I attempted to turn off the gradient for the stn.theta[0][1] and stn.theta[1][0] parts of the affine tranformation but this wasn’t possible due to them not being leaf nodes.
Instead I’ve created a tensor to prevent the shear parts of the affine transform from having any effect. Is this the correct approach?
T = torch.tensor
class AffineTransform(nn.Module):
    def __init__(self, scalew=1, scaleh=1, transX=0., transY=0.):
        super().__init__()
        self.theta = nn.Parameter(T([
            [scalew, 0, transX],
            [0, scaleh, transY]
        ], dtype=torch.float)[None])
        self._stop_shear = T([[1,0,1],[0,1,1]], requires_grad=False)
        
        
    def forward(self, x):
        theta = self.theta * self._stop_shear
        grid = F.affine_grid(theta, x.size(), align_corners=True)
        return F.grid_sample(x, grid, align_corners=True)