code: def test_replace_layers_in_module_swap_all_relu(a_model):
      mapping = (
          lambda mod: None
          if not isinstance(mod, nn.ReLU)
          else nn.Identity()
      )
      changed = replace_layers_in_module(a_model, mapping)
      assert changed
      assert not any(isinstance(m, nn.ReLU) for m in a_model.modules())
      assert any(isinstance(m, nn.Identity) for m in a_model.modules())
  
file_path: baal/tests/bayesian/common_test.py
imports: import pytest
  from torch import nn
  from baal.bayesian.common import replace_layers_in_module
label_desc: Use image node_id 3 for calling a global function with example usage: test_replace_layers_in_module_swap_all_relu(a_model) without return types