Slicing tensor breaks backprop

I have an input X that goes into an MLP which outputs a tensor xyz of shape torch.Size([4, 64, 1024]). I want to backprop to the input using only the first 3 columns of the last dimension. But when I slice xyz using xyz[:, :, 0:3] what happens is that autograd doesn’t backprop to the input X anymore. I also tried creating a linear layer and assign its weights manually from a tensor so it only selects the columns I want but Pytorch gives an error and it won’t work. Any ideas on what I should do?

Can you provide an script which shows that? It should work.