Solved by ezyang in post #2 
              
              
                I can’t think of a direct way to do this, but I guess if you’re training you could check which parameters actually got gradients and count only those.