search

f-dangel / curvlinops

scipy linear operators for the Hessian, Fisher/GGN, and more in PyTorch
https://curvlinops.readthedocs.io/en/latest/
MIT License
17 stars 5 forks source link

[OPT] Remove `for` loop to compute `grad_outputs` in EF #87

Closed f-dangel closed 6 months ago

f-dangel commented 6 months ago

The for-loop over data points in the current implementation can be a bottleneck for small networks. It will be better to parallelize the grad_output computation into a single call to torch.autograd.grad.