Fix DataLoader batch size mismatch across training phases + --train-iters now overrides all phases

- DatasetLoader.sample_target() recreates DataLoader when batch size changes
(Phase 1 pool building uses 256, Phase 2 NSF uses 128 — caused RuntimeError)
- --train-iters CLI flag now also overrides time_predictor iterations"

dataset_loader.py

@@ -140,7 +140,10 @@ class DatasetLoader:
 
     def sample_target(self, n: int, device: str = "cpu") -> torch.Tensor:
         if self.is_image:
-            if not hasattr(self, "_image_loader"):
+            # Recreate DataLoader if batch size changed (different training phases
+            # use different batch sizes, e.g. 256 for pool building, 128 for NSF)
+            if not hasattr(self, "_image_loader") or self._image_batch_size != n:
+                self._image_batch_size = n
                 self._image_loader = get_image_dataloader(
                     self.dataset_name, batch_size=n, train=True
                 )