code: def _consistent_dropout_mapping_fn(
      module: torch.nn.Module,
  ) -> Optional[nn.Module]:
      new_module: Optional[nn.Module] = None
      if isinstance(module, torch.nn.Dropout):
          new_module = ConsistentDropout(p=module.p)
      elif isinstance(module, torch.nn.Dropout2d):
          new_module = ConsistentDropout2d(p=module.p)
      return new_module
  
imports: import copy
  import warnings
  from typing import Optional
  import torch
  from torch import nn
  from torch.nn import functional
  from torch.nn.modules.dropout import _DropoutNd
  from baal.bayesian.common import replace_layers_in_module, _patching_wrapper, BayesianModule
label_desc: Use image node_id 3 for calling a global function with example usage: _consistent_dropout_mapping_fn(module) and returns: new_module