def override_assert_equal():
    r"""Override assertEqual to avoid triggering fp64 error on tensor comparison in test case"""

    def args_to_xpu(args):
        if torch.is_tensor(args) and args.is_xpu:
            return args.to("cpu")
        elif isinstance(args, (tuple, list)):
            args_list = list(args)
            for i, arg in enumerate(args_list):
                args_list[i] = args_to_xpu(arg)
            if isinstance(args, tuple):
                return tuple(args_list)
            else:
                return args_list
        else:
            return args

    def fp64_assert_equal_wrapper(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            args = args_to_xpu(args)
            return f(*args, **kwargs)

        return wrapper

    r"""
    In PyTorch design, `__allow_nonbracketed_mutation_flag` is a flag to forbid bare assignment
    to torch.backends.<cudnn|mkldnn>.enabled and friends when running test suite. This flag will
    be forced to set to False by function `disable_global_flags` which is defined in
    torch.testing._internal.common_utils when overriding TestCase.assertEqual. It may result in
    a runtime error on subsequent cudnn|mkldnn setting, if any. The function here is to override
    `disable_global_flags` with an empty one to keep the flag `__allow_nonbracketed_mutation_flag`
    from being changed.
    """

    def _disable_global_flags():
        pass

    torch.backends.disable_global_flags = _disable_global_flags
    from torch.testing._internal.common_utils import TestCase

    TestCase.assertEqual = fp64_assert_equal_wrapper(TestCase.assertEqual)
Override assertEqual to avoid triggering fp64 error on tensor comparison in test case
intel_extension_for_pytorch/xpu/overrides.py