Here is the problem broken down into a minumum working example:
    s = 100

    data = s*torch.ones(s)
    data.requires_grad=True

    target = torch.rand(s)
    target.requires_grad=True

    net = nn.Linear(s,10)
    optimizer = torch.optim.Adam(net.parameters(),1e-5)

    for i in range(s):
        optimizer.zero_grad()
        idx = net(data).clamp(0,s-1).long()
        loss = torch.index_select(target, 0, idx).mean()
        loss.backward()
        optimizer.step()

        print(loss)

The issue is that the linear layer doesn’t receive gradients and the reason is that idx.requires_grad=False due to the long conversion, which is required by index_select.