diff --git a/py/torch_tensorrt/dynamo/_refit.py b/py/torch_tensorrt/dynamo/_refit.py index 43fdecbf73..2b3ddce96e 100644 --- a/py/torch_tensorrt/dynamo/_refit.py +++ b/py/torch_tensorrt/dynamo/_refit.py @@ -37,7 +37,6 @@ from torch_tensorrt.dynamo.runtime._TRTEngine import TRTEngine from torch_tensorrt.dynamo.utils import ( check_module_output, - check_output_equal, get_torch_inputs, to_torch_device, to_torch_tensorrt_device, @@ -436,31 +435,15 @@ def refit_module_weights( if verify_output and arg_inputs is not None: new_gm.to(to_torch_device(settings.device)) # move to device for inference - # complex_graph_detection rewrites complex placeholders to real (view_as_real). - # The compiled TRT module handles complex→real internally, but the lowered - # PyTorch reference module (new_gm) expects real-unpacked inputs directly. - has_complex_inputs = any( - isinstance(x, torch.Tensor) and x.is_complex() for x in torch_inputs + # A complex placeholder stays complex in the lowered reference module; the + # unpacking to real happens inside the graph, so both modules take the same + # inputs the caller passed. + outputs_match = check_module_output( + new_module=new_gm, + refitted_module=compiled_module, + arg_inputs=torch_inputs, + kwarg_inputs=torch_kwarg_inputs, ) - if has_complex_inputs: - lowered_inputs = [ - ( - torch.view_as_real(x).contiguous() - if isinstance(x, torch.Tensor) and x.is_complex() - else x - ) - for x in torch_inputs - ] - trt_outputs = compiled_module(*torch_inputs) - ref_outputs = new_gm(*lowered_inputs, **torch_kwarg_inputs) - outputs_match = check_output_equal(trt_outputs, ref_outputs) - else: - outputs_match = check_module_output( - new_module=new_gm, - refitted_module=compiled_module, - arg_inputs=torch_inputs, - kwarg_inputs=torch_kwarg_inputs, - ) if outputs_match: logger.info("Refitting Succeed!") else: