It can vectorize over a for-loop, yes. Just bare in mind, you can’t mix torch.autograd.grad with torch.func, so you’ll need to re-write how your loss function is defined.
It can vectorize over a for-loop, yes. Just bare in mind, you can’t mix torch.autograd.grad with torch.func, so you’ll need to re-write how your loss function is defined.