code: def freeze_model(
      module: nn.Module,
      freeze_parameters: Optional[bool] = True,
      freeze_bn: Optional[bool] = True,
  ) -> nn.Module:
      """
      Change 'requires_grad' value for module and it's child modules and
      optionally freeze batchnorm modules.
      :param module: Module to change
      :param freeze_parameters: True to freeze parameters; False - to enable parameters optimization.
          If None - current state is not changed.
      :param freeze_bn: True to freeze batch norm; False - to enable BatchNorm updates.
          If None - current state is not changed.
      :return: None
      """
      bn_types = (
          nn.BatchNorm1d,
          nn.BatchNorm2d,
          nn.BatchNorm3d,
          nn.SyncBatchNorm,
      )
  
      if freeze_parameters is not None:
          for param in module.parameters():
              param.requires_grad = not freeze_parameters
  
      if freeze_bn is not None:
          if isinstance(module, bn_types):
              module.track_running_stats = not freeze_bn
  
          for m in module.modules():
              if isinstance(m, bn_types):
                  module.track_running_stats = not freeze_bn
  
      return module
  
file_path: pytorch-toolbelt/pytorch_toolbelt/optimization/functional.py
imports: import dataclasses
  import numbers
  from typing import Optional, Iterator, Dict, Union, List, Tuple, Mapping
  from torch import nn
  from pytorch_toolbelt.utils.distributed import get_rank, get_world_size, is_dist_avail_and_initialized
  from pytorch_toolbelt.utils.torch_utils import get_non_wrapped_model
label_desc: Use image node_id 6 for calling a global function with example usage: freeze_model(module, freeze_parameters, freeze_bn) and returns: module