Make test_imputation unit test fair by training the model with imputation for the same number of epochs as the model without imputation