How to find energy per layer using pytorch ?
Hi @coolcoder001! I am interested in the same issue. Did you found some way to do it?