Use ..._batch_size rather than per_device_..._batch_size

The former is identical for single-GPU and DDP, but has a higher batch size for DP (which is the expected behaviour).
UKPLab · Sep 10, 2024 · 5f91adb · 5f91adb
1 parent 6d7637a
commit 5f91adb
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/sentence_transformers/trainer.py b/sentence_transformers/trainer.py
@@ -571,7 +571,7 @@ def get_train_dataloader(self) -> DataLoader:
             batch_samplers = [
                 self.get_batch_sampler(
                     dataset,
-                    batch_size=self.args.per_device_train_batch_size,
+                    batch_size=self.args.train_batch_size,
                     drop_last=self.args.dataloader_drop_last,
                     valid_label_columns=data_collator.valid_label_columns,
                     generator=generator,
@@ -644,7 +644,7 @@ def get_eval_dataloader(self, eval_dataset: Dataset | None = None) -> DataLoader
             batch_samplers = [
                 self.get_batch_sampler(
                     dataset,
-                    batch_size=self.args.per_device_eval_batch_size,
+                    batch_size=self.args.eval_batch_size,
                     drop_last=self.args.dataloader_drop_last,
                     valid_label_columns=data_collator.valid_label_columns,
                     generator=generator,
@@ -708,7 +708,7 @@ def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
             batch_samplers = [
                 self.get_batch_sampler(
                     dataset,
-                    batch_size=self.args.per_device_eval_batch_size,
+                    batch_size=self.args.eval_batch_size,
                     drop_last=self.args.dataloader_drop_last,
                     valid_label_columns=data_collator.valid_label_columns,
                     generator=generator,