Adaptive pooling for diagh

f-dangel / backpack

BackPACK - a backpropagation package built on top of PyTorch which efficiently computes quantities other than the gradient.

https://backpack.pt/

MIT License

549 stars 55 forks source link

Closed fKunstner closed 10 months ago

fKunstner commented 10 months ago

Should fix #313

I think it was just a case of a missing module_ext.

Changes:

Added class DiagHAdaptiveAvgPoolNd(DiagHBaseModule)
Registered it in class DiagHessian(SecondOrderBackpropExtension), class BatchDiagHessian(SecondOrderBackpropExtension)
Added tests for the combination