[feature request] KFAC support for ResNet

Hi, thank you for developing and maintaining this library.

If I understand correctly, backpack 1.4 does not support ResNet with KFAC, right?

When I changed DiagGGNExact of the tutorial to KFAC (and changed AdaptiveAvgPooling to AvgPooling), it raised the following error:

~/.miniconda/.../backpack/extensions/backprop_extension.py in __get_module_extension(self, module)
     97             if self._fail_mode is FAIL_ERROR:
     98                 # PyTorch converts this Error into a RuntimeError for torch<1.7.0
---> 99                 raise NotImplementedError(
    100                     f"Extension saving to {self.savefield} "
    101                     "does not have an extension for "

NotImplementedError: Extension saving to kfac does not have an extension for Module <class 'backpack.custom_module.branching.SumModule'>

It would be great if you could support KFAC for ResNet or let me know some modifications necessary if possible. Thank you.

[f-dangel]

Hi,

thanks for giving the tutorial a try. You correctly tweaked it, but support for K-FAC with ResNets is indeed not implemented in 1.4. @jabader97 is working on this feature and I can post updates once it makes it into the development branch.

Best, Felix

[moskomule]

Thank you again!