Is there an implementation for kfac that can handle batchnorm? I found this implementation of kfac but this does not support batchnorm. I need some hint on how to extend it to batchnorm. I greatly appreciate it.

  
      
      GitHub
  
  
    

alecwangcq/KFAC-Pytorch 6

Pytorch implementation of KFAC and E-KFAC (Natural Gradient). - alecwangcq/KFAC-Pytorch