diff --git a/nnunetv2/run/run_training.py b/nnunetv2/run/run_training.py
index 73a5ce88d..fff2ecd2e 100644
--- a/nnunetv2/run/run_training.py
+++ b/nnunetv2/run/run_training.py
@@ -278,5 +278,8 @@ def run_training_entry():
 
 
 if __name__ == '__main__':
+    os.environ['OMP_NUM_THREADS'] = '1'
+    os.environ['MKL_NUM_THREADS'] = '1'
+    os.environ['OPENBLAS_NUM_THREADS'] = '1'
     # multiprocessing.set_start_method("spawn")
     run_training_entry()
diff --git a/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py b/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py
index fc79558ba..c5b144a23 100644
--- a/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py
+++ b/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py
@@ -367,8 +367,6 @@ def _set_batch_size_and_oversample(self):
 
             print("worker", my_rank, "oversample", oversample_percent)
             print("worker", my_rank, "batch_size", batch_size_per_GPU[my_rank])
-            # self.print_to_log_file("worker", my_rank, "oversample", oversample_percents[my_rank])
-            # self.print_to_log_file("worker", my_rank, "batch_size", batch_sizes[my_rank])
 
             self.batch_size = batch_size_per_GPU[my_rank]
             self.oversample_foreground_percent = oversample_percent
diff --git a/pyproject.toml b/pyproject.toml
index 77c14d8d1..d61925466 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,7 +49,8 @@ dependencies = [
     "matplotlib",
     "seaborn",
     "imagecodecs",
-    "yacs"
+    "yacs",
+    "batchgeneratorsv2"
 ]
 
 [project.urls]