Two things:



 Tomash:

                pred = F.softmax(output["age"], dim=-1).cpu().numpy()
                pred = (pred * np.arange(0, pred.size)).sum(axis=-1)
                preds.append(np.array([pred]))




What is happening in this block?
Have you trained your network for a sufficient number of epochs? Calculate the training f1 and mae to ensure the model is training as expected.