Your array has 3 dimensions and you’re trying to access it with 4.
You haven’t included all of the error message but I suppose it is train_x as that has 4 dimensions in your snippet.
So wherever that is assigned is where your problem might lie.
I would advise reading this advice