How to predict (multi) labeled datapoints?

for example if I have 5 points and each point has dimension of 3 and has one of 10 possible labels from 0 to 9.
Points

XTrain:
[[0.20861965 0.47901568 0.92075312], 
[0.96175914 0.70659989 0.82364516], 
[0.51805523 0.42727509 0.92545694], 
[0.4061363  0.55752676 0.56914541], 
[0.47859976 0.81323072 0.042954  ]]

Labels

y_true: [5 5 0 9 3] 

I convert the labels with to_categorical and get this:

y_true:
[[0. 0. 0. 0. 0. 1. 0. 0. 0. 0.], 
[0. 0. 0. 0. 0. 1. 0. 0. 0. 0.], 
[1. 0. 0. 0. 0. 0. 0. 0. 0. 0.], 
[0. 0. 0. 0. 0. 0. 0. 0. 0. 1.], 
[0. 0. 0. 1. 0. 0. 0. 0. 0. 0.]] 

Then I multiply (randomly initialised) theta with the points:

theta = np.random.normal(size=(3, 1))
a = np.matmul(XTrain, theta)

After that, apply activation function:

y_pred = tf.nn.sigmoid(a)

Now I want to calculate loss with (or should I use something else?)

loss = tf.keras.metrics.categorical_crossentropy(y_true, y_pred)

But the shape of y_pred is not correct and I got an error.

y_pred:
tf.Tensor(
[[0.69186984]
 [0.50125371]
 [0.4464451 ]
 [0.38257534]
 [0.62392589]], shape=(5, 1), dtype=float64)

I assume, this are the probabilities of each point beeing correctly predicted. But this is not what I want.
How should I calculate y_pred so it would have the correct format and shape like in the documentation of the loss for categorical_crossentropy loss? Example input for this loss from the documentation:

y_true = [[0, 1, 0], [0, 0, 1]]
y_pred = [[0.05, 0.95, 0], [0.1, 0.8, 0.1]]

y_true is like I have (categorical) but how do I get y_pred like this?