Problems with Flux NN regression

I am trying to make a regression with NN using Flux. I have separated the dataset in the train/test parts and implemented a simple neural network.

x_train, x_test, y_train, y_test = train_test_split(convert(Array, x), convert(Array, y), test_size=.3)

nndata = Flux.Data.DataLoader((x_train’, y_train), batchsize=30, shuffle=false)
model = Chain(
Dense(23, 40, relu),
Dense(40, 25, relu),
Dense(25, 1, identity), 
)
ps = Flux.params(model)
loss(x_train, y_train) = Flux.mse(model(x_train), y_train)
opt = ADAM()
using Flux: train!
for i in Array((1:50)’)
train!(loss, ps, nndata, opt)
end
y_pred = model(x_train’)

When I train the model, after some epochs, all the predicted value points to a single value, which is close to the mean of y_train.
Does anyone have a guess as on what is going on?
Thank you

Hi,
Not sure what the issue could be, but a starting point could be to look at how the loss is evolving during training. You can use callbacks to monitor this when using the train!() function.

It might also be helpful if you could describe what type of data you are dealing with. Do you have discrete or continuous labels? Is it a balanced dataset?