The validation code in test.py seems like not compatible with training model. Would you like to check it?