Trouble with importing tensorflow keras model into Barracuda.

So I’ve trained a tensorflow keras model on a dataset of digits. 0 to 9 to be exact. While testing in python it gives me 96% accuracy so I know it works.
.
Now I wanted to import it into unity using barracuda and the MLagents package (which uses tensorflow).
.
Step 1 was to convert my model.pb to a model.onnx
Step 2 was to put it in unity
Step 3 was giving it input (array of pixels in grayscale) and waiting for the output.
.
Now it should give me an output that is an array of floats. Each entry giving me a value from 0 to 1 describing how much it thinks it looks like the number of that index. So for the first entry, it would check if it’s similar to a 0, then 1 and so on.
.
When I test the whole process in unity, no matter what number i draw, it gives me around 0.9 for index 1 of the prediction array but for all the other entries it gives me a value very close to 0. Meaning it thinks it’s always number 1.
.
I am unsure what I am doing wrong and would gladly appreciate some help.
.

private void AnalyzeImage(Texture2D texture)
    {
        bool verbose = false;

        var additionalOutputs = new string[] { "StatefulPartitionedCall/sequential/dense/softmax" };
        var engine = WorkerFactory.CreateWorker(m_RuntimeModel, additionalOutputs, WorkerFactory.Device.GPU, verbose);

        texture = Resize(texture, 45, 45);

        //You can treat input pixels as 1 (grayscale), 3 (color) or 4 (color with alpha) channels
        var channelCount = 1;
        var input = new Tensor(texture, channelCount);

        engine.Execute(input);
        var prediction = engine.PeekOutput(/*"StatefulPartitionedCall/sequential/dense/softmax"*/);
        float[] predictionValues = prediction.AsFloats();
        foreach (float f in predictionValues)
        {
            Debug.Log(f);
        }

        float highestValue = 0;
        int index = 0;
        int highestValueIndex = 0;
        foreach (float i in predictionValues)
        {
            if (i > highestValue)
            {
                highestValue = i;
                highestValueIndex = index;
            }
            index += 1;
        }

        if (highestValue > 0.75f)
            m_UIText.text = "Recognized: " + m_ClassNames[highestValueIndex];
        else
            m_UIText.text = "Unable to Recognize";
    }

That’s the function I use to process data using the model.

Link to .onnx file download
Link to .pb file download

if you are testing it out in unity, the texture2D input has to be 45 by 45 and grayscale

take a look at this Github project, it’s about using Onnx models in Unity :