list -> generator

lsst · May 11, 2023 · 2cfc252 · 2cfc252
1 parent a790214
commit 2cfc252
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/python/lsst/meas/transiNet/rbTransiNetInterface.py b/python/lsst/meas/transiNet/rbTransiNetInterface.py
@@ -72,7 +72,7 @@ def init_model(self):
         self.model.eval()
 
     def input_to_batches(self, inputs, batchSize):
-        """Convert a list of inputs to batches of inputs.
+        """Convert a list of inputs to a generator of batches.
 
         Parameters
         ----------
@@ -81,11 +81,11 @@ def input_to_batches(self, inputs, batchSize):
 
         Returns
         -------
-        batches : `list` [`list` [`CutoutInputs`]]
-            List of batches of inputs.
+        batches : `generator`
+            Generator of batches of inputs.
         """
-        batches = [inputs[i:i + batchSize] for i in range(0, len(inputs), batchSize)]
-        return batches
+        for i in range(0, len(inputs), batchSize):
+            yield inputs[i:i + batchSize]
 
     def prepare_input(self, inputs):
         """