Allow num_batches to be None

alessiamarcolini · alessiamarcolini · commit 0fee00272e62 · 2022-07-11T15:46:03.000+02:00
diff --git a/test/test_utils.py b/test/test_utils.py
@@ -36,16 +36,16 @@ def default_dataloader(self, batch_size=None, num_samples=None, shuffle=True):
             self.ds, batch_size=batch_size, sampler=sampler
         )
 
-
+@pytest.mark.parametrize("num_batches", (3, None))
 @pytest.mark.parametrize("ds_type", ("map", "iterable"))
-def test_fit(ds_type):
+def test_fit(ds_type, num_batches):
     dm_ref = TensorDataModule(torch.randn(500, 5), ds_type=ds_type)
     d = torchdrift.detectors.KernelMMDDriftDetector()
     torchdrift.utils.fit(
         dm_ref.default_dataloader(batch_size=10),
         torch.nn.Identity(),
         [torch.nn.Identity(), d],
-        num_batches=3,
+        num_batches=num_batches,
         device="cpu",
     )
 
diff --git a/torchdrift/utils/fit.py b/torchdrift/utils/fit.py
@@ -43,9 +43,9 @@ def fit(
         total = None
 
     for i, b in enumerate(tqdm.tqdm(dl, total=total)):
-        if i >= num_batches:
+        if num_batches and i >= num_batches:
             break
-        
+
         if not isinstance(b, torch.Tensor):
             b = b[0]
         with torch.no_grad():
