Onnx export with dynamo using torch.cond for dynamic models

ionymikler: <h3><a name="p-461487-describe-the-bug-1" class="anchor" href="#p-461487-describe-the-bug-1"></a><img src="https://discuss.pytorch.org/images/emoji/apple/bug.png?v=12" title=":bug:" class="emoji" alt=":bug:" loading="lazy" width="20" height="20"> Describe the bug</h3>
<p><strong>Environment</strong><br>
pytorch.version = 2.5.1+cu124’</p>
<p><strong>Description</strong><br>
I am trying to implement a dummy example of a model whose forward method operations would depend on some intermediate calculation on the input. The final goal of this is to see if I can export such a model to ONNX.</p>
<p>The model in question is the following:</p>
<pre data-code-wrap="python"><code class="lang-python">class TwoLayerNetDynamic(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(TwoLayerNetDynamic, self).__init__()
        self.model_name = 'TwoLayerNetDynamic'

        self.fully_connected_1 = nn.Linear(input_size, hidden_size)
        self.early_exit_head_1 = nn.Linear(hidden_size, output_size)

        self.threshold = torch.tensor([0.0], dtype=torch.float32).to(DEVICE)
        self.last_exit = nn.Linear(hidden_size, output_size)

        self.training_exits = False

        print("TwoLayerNetDynamic initialized")

    def forward(self, x: torch.Tensor):
        mean = x.mean()
        
        # version 1 
        # if torch.gt(mean, self.threshold):
        #     x = self.early_exit_head_1(x)
        # else:
        #     x = self.last_exit(x)
        # x = torch.cat([x, mean.reshape_as(x)], dim=1)

        # version 2
        x = self.fully_connected_1(x)
        x = fc.cond(mean&gt;0.0,self.early_exit_head_1,self.last_exit,(x,))
        x = torch.cat([x, mean.reshape_as(x)], dim=1)       
        return x 
</code></pre>
<p>Using the commented version1 of <code>forward</code>, I used <code>torch.jit_script</code> as mentioned <a href="https://pytorch.org/docs/stable/generated/torch.jit.script.html#torch.jit.script" rel="noopener nofollow ugc">here</a>.<br>
I manage to export a onnx model that has the dynamic behavior I want doing the following:</p>
<pre data-code-wrap="python"><code class="lang-python">model = TwoLayerNetDynamic(input_size=1, hidden_size=3, output_size=1)

script_module = torch.jit.script(model)
    torch.onnx.export(
        model=script_module,
        args=_x,
        f=f"./models/onnx/{model.model_name}_scripting.onnx"
)
</code></pre>
<p>I at least think this is the case because when I see the exported model in Netron i see the flow-control statement captured:<br>
<img src="https://github.com/user-attachments/assets/e303119e-b29d-4b6e-8155-9f172aaf9d5d" alt="image" width="288" height="500"></p>
<p><strong>So, what is the problem?</strong><br>
When I try doing the same thing with the new API using Dynamo, it gives me an error.<br>
The procedure there is this:</p>
<pre data-code-wrap="python"><code class="lang-python">### Using TorchDynamo ###
filename = f""
onnx_filepath = f"./models/onnx/{model.model_name}_dynamo.onnx"
onnx_program:ONNXProgram = torch.onnx.export(
  model=model,
  args=(_x,),
  dynamo=True,
  report=True
)
onnx_program.save(onnx_filepath)
</code></pre>
<p>When doing this, I get a lenghty error, but also the markdown report for the export. Th error is:</p>
<pre><code class="lang-auto">Traceback (most recent call last):

  File "/home/iony/miniconda3/envs/lgvit/lib/python3.9/site-packages/torch/onnx/_internal/exporter/_core.py", line 553, in _add_nodes
    _handle_call_function_node_with_lowering(

  File "/home/iony/miniconda3/envs/lgvit/lib/python3.9/site-packages/torch/onnx/_internal/exporter/_core.py", line 444, in _handle_call_function_node_with_lowering
    raise _errors.DispatchError(

torch.onnx._internal.exporter._errors.DispatchError: No ONNX function found for &lt;torch._higher_order_ops.cond.CondOp object at 0x7f87bb7a7a30&gt;. Failure message: No decompositions registered for the real-valued input


The above exception was the direct cause of the following exception:


Traceback (most recent call last):

  File "/home/iony/miniconda3/envs/lgvit/lib/python3.9/site-packages/torch/onnx/_internal/exporter/_core.py", line 1134, in export
    onnx_program = _exported_program_to_onnx_program(

  File "/home/iony/miniconda3/envs/lgvit/lib/python3.9/site-packages/torch/onnx/_internal/exporter/_core.py", line 791, in _exported_program_to_onnx_program
    values = _add_nodes(exported_program, model, lower=lower, registry=registry)

  File "/home/iony/miniconda3/envs/lgvit/lib/python3.9/site-packages/torch/onnx/_internal/exporter/_core.py", line 565, in _add_nodes
    raise _errors.ConversionError(

torch.onnx._internal.exporter._errors.ConversionError: Error when translating node %cond : [num_users=1] = call_function[target=torch.ops.higher_order.cond](args = (%gt, %true_graph_0, %false_graph_0, (%addmm, %p_early_exit_head_1_bias, %p_early_exit_head_1_weight, %p_last_exit_bias, %p_last_exit_weight)), kwargs = {}). See the stack trace for more information.
</code></pre>
<p>And from what I can see the <code>torch.cond</code> function is not supported yet.<br>
Searching online I came across <a href="https://github.com/pytorch/pytorch/pull/137428/files#diff-e97696ebe6c71bffa8517a9dc03c10f1d3dcb3c036b2128c1699456c4346052d" rel="noopener nofollow ugc">this PR to pytorch</a>.<br>
So it looks like this High-level Operations (HOPs) are supported.</p>
<p>My question is, am I doing something wrong?<br>
Any insights?</p>
<h1><a name="p-461487-environment-2" class="anchor" href="#p-461487-environment-2"></a>Environment</h1>
<p>PyTorch version: 2.5.1+cu124<br>
Is debug build: False<br>
CUDA used to build PyTorch: 12.4<br>
ROCM used to build PyTorch: N/A</p>
<p>OS: Ubuntu 20.04.6 LTS (x86_64)<br>
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0<br>
Clang version: Could not collect<br>
CMake version: version 3.16.3<br>
Libc version: glibc-2.31</p>
<p>Python version: 3.9.20 (main, Oct  3 2024, 07:27:41)  [GCC 11.2.0] (64-bit runtime)<br>
Python platform: Linux-5.15.0-126-generic-x86_64-with-glibc2.31<br>
Is CUDA available: True<br>
CUDA runtime version: 12.4.131<br>
CUDA_MODULE_LOADING set to: LAZY</p>