PyTorch TransformerEncoder

Hi,

I am trying to convert a model containing PyTorch TransformerEncoder but I am getting this error with Barracuda 3.0.0

Exception: Must have input rank for /seqTransEncoder/layers.0/self_attn/Slice_output_0 in order to convert Reshape to NHWC
Unity.Barracuda.Compiler.Passes.NCHWToNHWCPass.<InstantiateRewriterNCHWToNHWC>b__4_0 (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:80)
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 <964670f15c6b47f9b8f9340732720473>:0)

My model is here

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        trans_enc_layer = nn.TransformerEncoderLayer(d_model=32,
                                                     nhead=4,
                                                     dim_feedforward=64)

        self.seqTransEncoder = nn.TransformerEncoder(trans_enc_layer,
                                                     num_layers=1)

    def forward(self, x):
        return self.seqTransEncoder(x)

model = Model()
x = torch.randn(7, 1, 32)
y = model(x)

output_file_name = 'model.onnx'
torch.onnx.export(
    model,
    x,
    output_file_name,
    training=torch.onnx.TrainingMode.EVAL,
)

In addition, the outputs I get in PyTorch and onnxruntime seems to be different. When I use

ort_session = ort.InferenceSession(output_file_name)
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)
np.testing.assert_allclose(to_numpy(y), ort_outs[0], rtol=1e-03, atol=1e-05)

I get the following

Not equal to tolerance rtol=0.001, atol=1e-05

Mismatched elements: 224 / 224 (100%)
Max absolute difference: 0.96960294
Max relative difference: 14.14618

Unity-Technologies / barracuda-release

PyTorch TransformerEncoder #330