def __init__(self, size, mode='bilinear'):
    super().__init__()
    self.mode = mode
    vectors = [torch.arange(0, s) for s in size]
    grids = torch.meshgrid(vectors)
    grid = torch.stack(grids)
    grid = torch.unsqueeze(grid, 0)
    grid = grid.type(torch.FloatTensor)
    self.register_buffer('grid', grid)
def __init__(self, size, mode='bilinear'):
    super().__init__()
    self.mode = mode
    vectors = [torch.arange(0, s) for s in size]
    grids = torch.meshgrid(vectors)
    grid = torch.stack(grids)
    grid = torch.unsqueeze(grid, 0)
    grid = grid.type(torch.FloatTensor)
    self.register_buffer('grid', grid)