Prepare for label noise experiments

wearepal · Oct 2, 2023 · c725c4b · c725c4b
1 parent 68394b9
commit c725c4b
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/algs/adv/supmatch.py b/src/algs/adv/supmatch.py
@@ -33,7 +33,11 @@ def _get_data_iterators(self, dm: DataModule) -> tuple[IterTr, IterDep]:
         dl_tr = dm.train_dataloader(balance=True)
         # The batch size needs to be consistent for the aggregation layer in the setwise neural
         # discriminator
-        dl_dep = dm.deployment_dataloader(batch_size=dm.batch_size_tr)
+        dl_dep = dm.deployment_dataloader(
+            batch_size=dl_tr.batch_sampler.batch_size
+            if dm.deployment_ids is None
+            else dm.batch_size_tr
+        )
         return iter(dl_tr), iter(dl_dep)
 
     @override