Solved by sbelharbi in post #2 
              
              
                IntermediateLayerGetter tracks all the parameters before the last layer. so, this is fine. 
import sys
import torch
import torchvision
from torchvision.models._utils import IntermediateLayerGetter


def count_params(model):
    """
    Count the number of parameters.
    :param model:
    :return:
 …