search

aleximmer / Laplace

Laplace approximations for Deep Learning.
https://aleximmer.github.io/Laplace
MIT License
432 stars 61 forks source link

`torch.func` jacobian OOM #186

Open wiseodd opened 1 month ago

wiseodd commented 1 month ago

Minimum example to get this error: https://gist.github.com/wiseodd/b29973cd96f96f3af620ca131571eaa4

Fix:

use AsdlGGN backend and then in _glm_predictive_distribution (in baselaplace.py), use:

Js, f_mu = self.backend.jacobians(X)

instead of:

Js, f_mu = self.backend.functorch_jacobians(X)
wiseodd commented 1 month ago

I think this is because: https://github.com/pytorch/functorch/issues/1058