After updating to PyTorch0.4, the bilinear sampling in the following code doesn’t give me correct gradients anymore (basically the code does exact same input as torch.nn.functional.grid_sample() and does the same functionality. I specified where the discrepancy happens between PyTorch04 and PyTorch03 putting FIXME comment in there). I confirmed that the exact same code works in PyTorch0.3.
Though PyTorch0.4 provides torch.nn.functional.grid_sample and it gives me correct gradients, I want to get the custom bilinear sampling working as I have a couple of different cases where I use the similar bilinear sampling scheme.
I really appreciate if anyone can help.
class SpatialTransformLayer(nn.Module):
    def __init__(self):
        super(SpatialTransformLayer, self).__init__()

    def _interpolate(self, im, xy):
        '''
        args:
            im: input image (B, C, H, W)
            xy: sampling coordinates (B, 2(x,y), outH, outW)

        return:
            (B, C, outH, outW) tensor
        '''

        B, C, H, W = im.size()
        _, _, out_H, out_W = xy.size()

        x, y = torch.chunk(xy, 2, dim=1)
        x = x.contiguous().view(-1)
        y = y.contiguous().view(-1)
        # print(x.size(), y.size())

        x0 = torch.floor(x).int()
        x1 = x0 + 1
        y0 = torch.floor(y).int()
        y1 = y0 + 1

        max_x = W - 1
        max_y = H - 1

        x0_clamp = torch.clamp(x0, 0, max_x)
        x1_clamp = torch.clamp(x1, 0, max_x)
        y0_clamp = torch.clamp(y0, 0, max_y)
        y1_clamp = torch.clamp(y1, 0, max_y)

        dim2 = W
        dim1 = W * H

        base = None
        if im.is_cuda:
            base = Variable((dim1 * torch.arange(B).int()).view(B, 1).expand(B, out_H * out_W).contiguous().view(-1).cuda())
        else:
            base = Variable((dim1 * torch.arange(B).int()).view(B, 1).expand(B, out_H * out_W).contiguous().view(-1))

        base_y0 = base + y0_clamp * dim2
        base_y1 = base + y1_clamp * dim2

        idx_y0_x0 = base_y0 + x0_clamp
        idx_y0_x1 = base_y0 + x1_clamp
        idx_y1_x0 = base_y1 + x0_clamp
        idx_y1_x1 = base_y1 + x1_clamp

        # (B,C,H,W) -> (B,H,W,C)
        im_flat = im.permute(0,2,3,1).contiguous().view(-1, C)
        i_y0_x0 = torch.gather(im_flat, 0, idx_y0_x0.unsqueeze(1).expand(-1,C).long())
        i_y0_x1 = torch.gather(im_flat, 0, idx_y0_x1.unsqueeze(1).expand(-1,C).long())
        i_y1_x0 = torch.gather(im_flat, 0, idx_y1_x0.unsqueeze(1).expand(-1,C).long())
        i_y1_x1 = torch.gather(im_flat, 0, idx_y1_x1.unsqueeze(1).expand(-1,C).long())
        
        # Check the out-of-boundary case.
        x0_valid = (x0 <= max_x) & (x0 >= 0)
        x1_valid = (x1 <= max_x) & (x1 >= 0)
        y0_valid = (y0 <= max_y) & (y0 >= 0)
        y1_valid = (y1 <= max_y) & (y1 >= 0)

        x0 = x0.float()
        x1 = x1.float()
        y0 = y0.float()
        y1 = y1.float()

        # FIXME: gradients through x, y here get broken
        w_y0_x0 = ((x1 - x) * (y1 - y) * (x1_valid * y1_valid).float()).unsqueeze(1)
        w_y0_x1 = ((x - x0) * (y1 - y) * (x0_valid * y1_valid).float()).unsqueeze(1)
        w_y1_x0 = ((x1 - x) * (y - y0) * (x1_valid * y0_valid).float()).unsqueeze(1)
        w_y1_x1 = ((x - x0) * (y - y0) * (x0_valid * y0_valid).float()).unsqueeze(1)

        output = w_y0_x0 * i_y0_x0 + w_y0_x1 * i_y0_x1 + w_y1_x0 * i_y1_x0 + w_y1_x1 * i_y1_x1
        mask = (x0_valid * x1_valid * y0_valid * y1_valid).float()

        # (B,H,W,C) -> (B,C,H,W)
        return output.view(B, out_H, out_W, C).permute(0,3,1,2).contiguous(), mask.view(B, out_H, out_W, 1).permute(0,3,1,2).contiguous()

    def forward(self, im, xy):
        '''
        args:
            im: (B, C, H, W) tensor
            xy: (B, 2(xy), outH, outW) tensor
        return:
            (B, C, outH, outW) tensor
        '''
        B, C, H, W = im.size()

        im, mask = self._interpolate(im, xy)

        return im, mask