I created a simple model using Torch
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        
    def forward(self, x, y):
        if self.training:
            return torch.matmul(x, y)
        return torch.add(x, y)

tensor1 = torch.Tensor([4])
tensor2 = torch.Tensor([3])
model = MyModel()
torch.onnx.export(model, (tensor1,tensor2), 'train_model.onnx', input_names=["x","y"], output_names=["z"], training=2)

In this, the exported ONNX file has only one Add node and does not include the Matmul node. I require to add both the nodes in the ONNX file and be able to switch between them while training & Inference.
Thanks.