def forward(self, src, flow):
    new_locs = (self.grid + flow)
    shape = flow.shape[2:]
    for i in range(len(shape)):
        new_locs[:, i, ...] = (2 * ((new_locs[:, i, ...] / (shape[i] - 1)) - 0.5))
    if (len(shape) == 2):
        new_locs = new_locs.permute(0, 2, 3, 1)
        new_locs = new_locs[(..., [1, 0])]
    elif (len(shape) == 3):
        new_locs = new_locs.permute(0, 2, 3, 4, 1)
        new_locs = new_locs[(..., [2, 1, 0])]
    return nnf.grid_sample(src, new_locs, align_corners=True, mode=self.mode)
def forward(self, src, flow):
    new_locs = (self.grid + flow)
    shape = flow.shape[2:]
    for i in range(len(shape)):
        new_locs[:, i, ...] = (2 * ((new_locs[:, i, ...] / (shape[i] - 1)) - 0.5))
    if (len(shape) == 2):
        new_locs = new_locs.permute(0, 2, 3, 1)
        new_locs = new_locs[(..., [1, 0])]
    elif (len(shape) == 3):
        new_locs = new_locs.permute(0, 2, 3, 4, 1)
        new_locs = new_locs[(..., [2, 1, 0])]
    return nnf.grid_sample(src, new_locs, align_corners=True, mode=self.mode)