diff --git a/test/test_imputation.py b/test/test_imputation.py
index 9dfea81f7e29c01ecb907a7ea99a6ceee8f0c936..e73894aa2df8fcc9693b4f4d23d97e7fc5226c20 100644
--- a/test/test_imputation.py
+++ b/test/test_imputation.py
@@ -96,7 +96,7 @@ def test_performance_improves_for_A0205_with_pretraining():
         X_pretrain=X_index_imputed,
         ic50_pretrain=ic50_imputed,
         sample_weights_pretrain=sample_weights_imputed,
-        n_training_epochs=500)
+        n_training_epochs=200)
 
     ic50_pred_with_imputation = \
         predictor_with_imputation.predict_ic50_for_kmer_encoded_array(X_index)