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