From 2373c0fc24659e14de1a2dd62b69ee76226f0efc Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <shoumikhin@gmail.com>
Date: Tue, 25 Aug 2026 19:15:14 -0700
Subject: [PATCH] fix(dynamo): stop unpacking complex inputs before the refit
 output check

Three refit tests fail on main:

  test_complex_buffer_refit
  test_dual_complex_buffer_refit
  test_complex_buffer_with_real_param_refit

  TypeError: view_as_real is only supported for complex tensors

After refitting, `refit_module_weights(verify_output=True)` runs the lowered
PyTorch module and the compiled module on the same inputs and compares the
results. For a model with a complex input it first converted that input to a
real tensor of shape (..., 2), because the lowering pass used to rewrite the
complex placeholder to a real one.

It does not do that any more. The placeholder stays complex and the graph
starts with its own `view_as_real`, so the reference module was being handed a
real tensor and then asked to unpack it again, which is what raises.

Pass the caller's inputs to both modules, which is what the non-complex path
already did, so the two paths become one.
---
 py/torch_tensorrt/dynamo/_refit.py | 33 ++++++++----------------------
 1 file changed, 8 insertions(+), 25 deletions(-)

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: