Hi,
I am not sure but I think the problem is that you ignore strides when you index your tensors in the kernel.
Does changing you CoolFunc that way helps ?
    @staticmethod
    def backward(ctx, grad):
        grad = grad.clone()
        dx = _C.cool_backward(grad,
                              grad.shape[0], grad.shape[1],
                              grad.shape[2], grad.shape[3])
        return dx