coral-ordinal consistently predicts only 0 in classification task

[RoyBeckerKristal]

Hi there,
for a pet project of nn-based classifier of an ordinal variable, I tried to use coral-ordinal, and it kept predicted only 0 instead of other values in the range (I didn't experience this problem when using standard multiclass-related parameters).
In an attempt to isolate the problem, I wrote the following simple code just to test whether I'm able to get coral-ordinal to work. It creates a data set of 1000 items, each with three randomized independent variables 0<x1,x2,x3<1, and a randomized dependent integer variable 0 <= y <= 4. For whatever reason, the predictions are always 0. Any idea why?
Thanks,
Roy.

import tensorflow as tf
from tensorflow import keras
import numpy as np
from sklearn.model_selection import train_test_split
import coral_ordinal as coral
import random
from sklearn import metrics
import pandas
from keras import models,layers,optimizers

df = pandas.DataFrame()
df.insert(len(df.columns),"x1", np.random.rand(1000))
df.insert(len(df.columns),"x2", np.random.rand(1000))
df.insert(len(df.columns),"x3", np.random.rand(1000))
y = np.random.randint(5,size=1000)

X_train, X_test, y_train, y_test = train_test_split(df,y,test_size=0.2,random_state=41)

nn = keras.models.Sequential()
nn.add(layers.Dense(7, input_shape=(3,), activation="relu"))
nn.add(coral.CoralOrdinal(num_classes=5))
print(nn.summary())
loss_fn = coral.OrdinalCrossEntropy(num_classes=5)
optim_fn = optimizers.legacy.Adam()
nn.compile(optimizer=optim_fn,loss=loss_fn,metrics=[coral.MeanAbsoluteErrorLabels()])
history = nn.fit(X_train,y_train,epochs=100,validation_split=0.2)
test_pred = nn.predict(X_test)
_ ,test_acc, *dummy = nn.evaluate(X_test,y_test)
print(test_acc)
y_pred = np.array([np.argmax(x) for x in test_pred])

print("y_test y_pred")
for i in range(len(y_test)):
print(y_test[i],y_pred[i])