fixbug: batch_size

caopulan · Jan 3, 2023 · 4333561 · 4333561
1 parent 92624f7
commit 4333561
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/training/encoder_trainer.py b/training/encoder_trainer.py
@@ -125,7 +125,7 @@ def configure_datasets(self):
         train_batch_size = self.opts.batch_size // self.opts.gpu_num
         test_batch_size = self.opts.test_batch_size // self.opts.gpu_num
         assert self.opts.batch_size == train_batch_size * self.opts.gpu_num, 'Train batch size is not a multiple of gpu num.'
-        assert self.opts.batch_size == test_batch_size * self.opts.gpu_num, 'Test batch size is not a multiple of gpu num.'
+        assert self.opts.test_batch_size == test_batch_size * self.opts.gpu_num, 'Test batch size is not a multiple of gpu num.'
         if self.opts.dist:
             train_sampler = DistributedSampler(
                 self.train_dataset,