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): 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!