Hi,
The problem here is that your function takes a single Tensor as input. The autograd works with Tensors as a unit and so it does not recognize X[:,0] as being a subset of X, but as a different Tensor. So for him, this new Tensor was not used during the backward.
You will have to either give each entry separately as a different tensor (and use torch.cat in your forward to recreate a single Tensor) to be able to ask gradients for a single one.
Otherwise, you could ask the gradients for X and then slice the gradients to get the part you want.