code: def _caching_mapping_fn(
      module: torch.nn.Module,
  ) -> Optional[nn.Module]:
      new_module: Optional[nn.Module] = None
      # Could add more
      if isinstance(module, (nn.Linear, nn.Conv2d)):
          new_module = LRUCacheModule(module=module)
      return new_module
  
file_path: baal/baal/bayesian/caching_utils.py
imports: from typing import Optional
  import torch
  from torch import nn, Tensor
  from baal.bayesian.common import BayesianModule, _patching_wrapper
label_desc: Use image node_id 1 for calling a global function with example usage: _caching_mapping_fn(module) and returns: new_module