In addition the accuracy I reach is lower ( 0.5555555820465088) and the Confussion matrix:
nx.Tensor<
s64[3][3]
EXLA.Backend<host:0, 0.2777305747.1958084637.177673>
[
[11, 0, 0],
[6, 5, 3],
[7, 0, 4]
]
The dimension of the model is f32[3][13] instead of f32[13][3].
The prediction function also needs the model as input parameter.