Log the warning only for the first on-the-fly sampler.

PiperOrigin-RevId: 676817928
google-deepmind · Sep 22, 2024 · cfa5954 · cfa5954
1 parent 9a5c2c0
commit cfa5954
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/clrs/_src/samplers.py b/clrs/_src/samplers.py
@@ -106,7 +106,9 @@ def __init__(
     self._track_max_steps = track_max_steps
 
     if num_samples < 0:
-      logging.warning('Sampling dataset on-the-fly, unlimited samples.')
+      logging.log_first_n(
+          logging.WARNING, 'Sampling dataset on-the-fly, unlimited samples.', 1
+      )
       if track_max_steps:
         # Get an initial estimate of max hint length
         self.max_steps = -1