diff --git a/code/train.py b/code/train.py
index 8e48e14009dff51ed92d7baba7b49760146347a9..ff0e9e0d4bd99e41421d19eb102d060a8b6b625d 100644
--- a/code/train.py
+++ b/code/train.py
@@ -141,7 +141,7 @@ def extracted_forest_size_job(extracted_forest_size_job_pb, extracted_forest_siz
         sub_models_dir_files = os.listdir(sub_models_dir)
         for file_name in sub_models_dir_files:
             if '.pickle' != os.path.splitext(file_name)[1]:
-                return
+                continue
             else:
                 already_exists = os.path.getsize(os.path.join(sub_models_dir, file_name)) > 0
                 break