train.py Data Dimension Check

[romankouz]

In train.py, we need to update how we check that the data dimension is correct because it only works for autoencoders.

[romankouz]

needed for predict.py as well