Explicitly check num_batches is not None

alessiamarcolini · lantiga · web-flow · commit 03a40d327dea · 2022-08-25T00:07:56.000+02:00
Co-authored-by: Luca Antiga &lt;luca.antiga@gmail.com&gt;
diff --git a/torchdrift/utils/fit.py b/torchdrift/utils/fit.py
@@ -43,7 +43,7 @@ def fit(
         total = None
 
     for i, b in enumerate(tqdm.tqdm(dl, total=total)):
-        if num_batches and i >= num_batches:
+        if num_batches is not None and i >= num_batches:
             break
 
         if not isinstance(b, torch.Tensor):
