From b0b699a01fdc23eff6332bf5d4b5125dbff67a26 Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Thu, 26 Dec 2024 18:45:56 +0000 Subject: [PATCH] Add unit test to test that isinstance(...) behaves as expected with custom module types. --- .../custom_modules/test_all_custom_modules.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/test_all_custom_modules.py b/tests/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/test_all_custom_modules.py index 2014f1896a..66aa356d4a 100644 --- a/tests/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/test_all_custom_modules.py +++ b/tests/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/test_all_custom_modules.py @@ -74,6 +74,17 @@ def layer_to_device_via_state_dict(layer: torch.nn.Module, device: str): layer.load_state_dict(state_dict, assign=True) +@parameterize_all_layer_types +def test_isinstance(orig_layer: torch.nn.Module, layer_input: torch.Tensor, supports_cpu_inference: bool): + """Test that isinstance() and type() behave as expected after wrapping a layer in a custom layer.""" + orig_type = type(orig_layer) + + apply_custom_layers_to_model(orig_layer) + + assert isinstance(orig_layer, orig_type) + assert type(orig_layer) is not orig_type + + @parameterize_all_devices @parameterize_all_layer_types def test_state_dict(device: str, orig_layer: torch.nn.Module, layer_input: torch.Tensor, supports_cpu_inference: bool):