Method jit_test_assert

tests/attr/test_jit.py:142–217

(self: BaseTest)

        """

        def jit_test_assert(self: BaseTest) -> None:
            model_1 = model
            attr_args = args
            if (
                mode is JITCompareMode.data_parallel_jit_trace
                or JITCompareMode.data_parallel_jit_script
            ):
                if not torch.cuda.is_available() or torch.cuda.device_count() == 0:
                    raise unittest.SkipTest(
                        "Skipping GPU test since CUDA not available."
                    )
                # Construct cuda_args, moving all tensor inputs in args to CUDA device
                cuda_args = {}
                for key in args:
                    if isinstance(args[key], Tensor):
                        cuda_args[key] = args[key].cuda()
                    elif isinstance(args[key], tuple):
                        cuda_args[key] = tuple(
                            elem.cuda() if isinstance(elem, Tensor) else elem
                            for elem in args[key]
                        )
                    else:
                        cuda_args[key] = args[key]
                attr_args = cuda_args
                model_1 = model_1.cuda()

            # Initialize models based on JITCompareMode
            if (
                mode is JITCompareMode.cpu_jit_script
                or JITCompareMode.data_parallel_jit_script
            ):
                model_2 = torch.jit.script(model_1)  # type: ignore
            elif (
                mode is JITCompareMode.cpu_jit_trace
                or JITCompareMode.data_parallel_jit_trace
            ):
                formatted_inputs = _format_tensor_into_tuples(args["inputs"])
                additional_args: Tuple[Any, ...] = (
                    _format_additional_forward_args(args["additional_forward_args"])
                    if "additional_forward_args" in args
                    and args["additional_forward_args"] is not None
                    else ()
                )
                if formatted_inputs is not None:
                    all_inps = formatted_inputs + additional_args
                else:
                    all_inps = additional_args
                model_2 = torch.jit.trace(model_1, all_inps)  # type: ignore
            else:
                raise AssertionError("JIT compare mode type is not valid.")

            attr_method_1 = algorithm(model_1)
            attr_method_2 = algorithm(model_2)

            if noise_tunnel:
                attr_method_1 = NoiseTunnel(attr_method_1)
                attr_method_2 = NoiseTunnel(attr_method_2)
            if attr_method_1.has_convergence_delta():