code: def model_numel(
      model, param_dims=[2, 4], param_types=["weight", "bias"]
  ):
      """Count the number elements in a model's parameter tensors"""
      total_numel = 0
      for name, param in model.state_dict().items():
          # Extract just the actual parameter's name, which in this context we treat as its "type"
          if param.dim() in param_dims and any(
              type in name for type in param_types
          ):
              total_numel += torch.numel(param)
      return total_numel
  
imports: import argparse
  from collections import OrderedDict
  from copy import deepcopy
  import logging
  import operator
  import random
  import numpy
  import torch
  import torch.nn
  import torch.backends.cudnn
  import yaml
  import inspect
label_desc: Use image node_id 31 for calling a global function with example usage: model_numel(model, param_dims, param_types) and returns: total_numel