Okaay now I get it.
Please observe the following snippet:
import torch
torch.manual_seed(10)
a =torch.Tensor(list(range(0,5))).requires_grad_()
b =torch.Tensor(list(range(5,10))).requires_grad_()

optim = torch.optim.SGD([a,b],lr=1)
def run(c):
    optim.zero_grad()
    c.backward()
    optim.step


print(a)
print(a.sum())
print(b)
print(b.sum())

c = torch.max(a.sum(),b.sum())
run(c)
print(a.grad)
print(b.grad)
print("grads doesn't flow through a as b is maximum, this is expected ")
print("Now let's suppose that you want to prune first 2 elements of b")

c = torch.max(a.sum(),b[2:].sum())

run(c)
print(a.grad)
print(b.grad)
print("As you can see gradients doesnt flow through b[0:1]")
print("same effect happens if you multiply the tensor element by 0, as derivative of constant*variable = constant*D_variable")

c = torch.max(a.sum(),(b*torch.Tensor([0,0,1,1,1])).sum())

run(c)
print(a.grad)
print(b.grad)

tensor([0., 1., 2., 3., 4.], requires_grad=True)
tensor(10., grad_fn=<SumBackward0>)
tensor([5., 6., 7., 8., 9.], requires_grad=True)
tensor(35., grad_fn=<SumBackward0>)
tensor([0., 0., 0., 0., 0.])
tensor([1., 1., 1., 1., 1.])
grads doesn't flow through a as b is maximum, this is expected 
Now let's suppose that you want to prune first 2 elements of b
tensor([0., 0., 0., 0., 0.])
tensor([0., 0., 1., 1., 1.])
As you can see gradients doesnt flow through b[0:1]
same effect happens if you multiply the tensor element by 0, as derivative of constant*variable = constant*D_variable
tensor([0., 0., 0., 0., 0.])
tensor([0., 0., 1., 1., 1.])

You just need to play with masking and torch.max to accomplish whatever