Problem with dynamic input size

I am trying to import an ONNX model that I exported from PyTorch, however importing only works when not setting the dynamic_axes parameter in PyTorch's torch.onnx.export() function. If I do make the input or output size dynamic, I get the following error when trying to import the .onnx file:

Exception: Must have input rank for 208 in order to convert axis for NHWC op
Unity.Barracuda.Compiler.Passes.NCHWToNHWCPass.Upsample (Unity.Barracuda.Layer layer, Unity.Barracuda.ModelBuilder net) (at Library/PackageCache/com.unity.barracuda@3.0.0/Barracuda/Runtime/Core/Compiler/Passes/NCHWToNHWC/RewriterNCHWToNHWC.cs:515)
Unity.Barracuda.Compiler.Passes.NCHWToNHWCPass.Rewrite (Unity.Barracuda.Model& model) (at Library/PackageCache/com.unity.barracuda@3.0.0/Barracuda/Runtime/Core/Compiler/Passes/NCHWToNHWCPass.cs:139)
Unity.Barracuda.Compiler.Passes.NCHWToNHWCPass.Run (Unity.Barracuda.Model& model) (at Library/PackageCache/com.unity.barracuda@3.0.0/Barracuda/Runtime/Core/Compiler/Passes/NCHWToNHWCPass.cs:39)
Unity.Barracuda.Compiler.Passes.IntermediateToRunnableNHWCPass.Run (Unity.Barracuda.Model& model) (at Library/PackageCache/com.unity.barracuda@3.0.0/Barracuda/Runtime/Core/Compiler/Passes/IntermediateToRunnableNHWCPass.cs:38)
Unity.Barracuda.ONNX.ONNXModelConverter.Convert (Google.Protobuf.CodedInputStream inputStream) (at Library/PackageCache/com.unity.barracuda@3.0.0/Barracuda/Runtime/ONNX/ONNXModelConverter.cs:188)
Unity.Barracuda.ONNX.ONNXModelConverter.Convert (System.String filePath) (at Library/PackageCache/com.unity.barracuda@3.0.0/Barracuda/Runtime/ONNX/ONNXModelConverter.cs:98)
Unity.Barracuda.ONNXModelImporter.OnImportAsset (UnityEditor.AssetImporters.AssetImportContext ctx) (at Library/PackageCache/com.unity.barracuda@3.0.0/Barracuda/Editor/ONNXModelImporter.cs:65)
UnityEditor.AssetImporters.ScriptedImporter.GenerateAssetData (UnityEditor.AssetImporters.AssetImportContext ctx) (at <ac4382394009481fa61779d0ab47a96f>:0)
UnityEditorInternal.InternalEditorUtility:ProjectWindowDrag(HierarchyProperty, Boolean)
UnityEngine.GUIUtility:ProcessEvent(Int32, IntPtr, Boolean&)

Now I'm wondering what I'm doing wrong, since I thought Barracuda supports dynamic input sizes and shapes.

Some more possibly relevant info: Barracuda Version: 3.0.0 (according to Package Manager) Unity Editor Version: 2021.2.10f ONNX Opset: 9 The code I used to export the model from PyTorch (excluding irrelevant sections):

dynamic_axes_mapping = {'input_img' : {2 : 'iheight', 3 : 'iwidth'},
                        'output_img' : {2 : 'oheight', 3 : 'owidth'}}

# Export the model
torch.onnx.export(model,                        # model being run
    dummy_input,                                # model input
    save_path,                                  # where to save the model
    export_params = True,                       # store the trained parameter weights inside the model file
    opset_version = 9,                          # the ONNX version to export the model to
    do_constant_folding = True,                 # whether to execute constant folding for optimization
    input_names = ['input_img'],                # the model's input names
    output_names = ['output_img'],              # the model's output names
    dynamic_axes = dynamic_axes_mapping         # dynamic input size/shape
)

The input and output are both Tensors in the NCHW format.

I would really appreciate help and didn't want to report a bug without knowing whether the error was on my end. My research on the forums and past discussions have not yielded great results either.

Thank you in advance!

Unity-Technologies / barracuda-release

Problem with dynamic input size #268