code: def count_parameters(module, only_trainable=False):
      s = sum(
          p.numel()
          for p in module.parameters(recurse=False)
          if not only_trainable or p.requires_grad
      )
      if isinstance(module, PathSamplingLayerChoice):
          s += sum(
              count_parameters(m) for m in module.sampled_choices()
          )
      else:
          s += sum(count_parameters(m) for m in module.children())
      return s
  
imports: import logging
  from collections import OrderedDict
  import time
  import numpy
  import torch
  import nni.retiarii.nn.pytorch
  from nni.nas.pytorch.mutables import Mutable, InputChoice, LayerChoice
label_desc: Use image node_id 2 for calling a global function with example usage: count_parameters(module, only_trainable) and returns: s