code: def test_caching(my_model):
      x = torch.rand(10)
  
      # No Caching
      my_model(x)
      assert LinearMocked.call_count == 4
      my_model(x)
      assert LinearMocked.call_count == 8
  
      with MCCachingModule(my_model) as model:
          model(x)
          assert LinearMocked.call_count == 12
          model(x)
          assert LinearMocked.call_count == 12
  
      # No Caching
      my_model(x)
      assert LinearMocked.call_count == 16
      my_model(x)
      assert LinearMocked.call_count == 20
  
file_path: baal/tests/bayesian/test_caching.py
imports: import warnings
  import pytest
  import torch
  from torch.nn import Sequential, Linear
  from baal import ModelWrapper
  from baal.bayesian.caching_utils import MCCachingModule
label_desc: Use image node_id 2 for calling a global function with example usage: test_caching(my_model) without return types