Solved.
Wrong typo in the loss function.
loss(x,y) = logitcrossentropy(model(x),y)
into:
loss(x,y) = logitcrossentropy(mean(model(x).nf, dims = 2),y)
Solved.
Wrong typo in the loss function.
loss(x,y) = logitcrossentropy(model(x),y)
into:
loss(x,y) = logitcrossentropy(mean(model(x).nf, dims = 2),y)