Given a trained model (M), I’m interested in computing the utility of new (unseen) examples in a pool (for an active learning task).  For this, I need to compute the magnitude of the gradient when M is trained on each new example.  In code, it is something like:
losses, grads = [], []
for i in range(X_pool.shape[0]):
    pred = model(X_pool[i:i+1])
    loss = loss_func(pred, y_pool[i:i+1])

    model.zero_grad()
    loss.backward()

    losses.append(loss)
    grads.append(layer.weight.grad.norm())

However, this is quite slow when there is a large pool of examples, especially since this will be an inner loop in my scenario.  How can this code be improved for efficiency?
I appreciate any suggestions or pointers.  Thanks!