pytorch-model.py

IN_WH = 512
GRID_WH = 256

class TestModel(nn.Module):

    def forward(self, x, grid):
        grid_resized = self.resize_grid(grid)
        return F.grid_sample(
            x, grid_resized
        )

    def resize_grid(self, grid):
        # [1, GRID_WH, GRID_WH, 2] => [1, 2, GRID_WH, GRID_WH]
        grid_resized = grid.permute(0, 3, 1, 2)
        # [1, 2, GRID_WH, GRID_WH] => [1, 2, IN_WH, IN_WH]
        grid_resized = F.interpolate(
            grid_resized, 
            size=(IN_WH, IN_WH), 
            mode='nearest'
        )
        # [1, 2, IN_WH, IN_WH] => [1, IN_WH, IN_WH, 2]
        grid_resized = grid_resized.permute(0, 2, 3, 1)
        return grid_resized