Solved by ptrblck in post #2 
              
              
                Yes, indexing the model output and target should work and the gradient would be backpropagated to the selected values: 
output = torch.randn(10, 1, requires_grad=True)
target = torch.randint(0, 2, (10, 1)).float()
criterion = nn.BCEWithLogitsLoss()

batch_idx = torch.tensor([1, 3, 5, 7])
loss = crit…