def forward(self, theta):
    warped_grid = self.apply_transformation(theta, torch.cat((self.grid_X, self.grid_Y), 3))
    return warped_grid
def forward(self, theta):
    warped_grid = self.apply_transformation(theta, torch.cat((self.grid_X, self.grid_Y), 3))
    return warped_grid