Another bugfix

csxeba · Mar 4, 2020 · 0b0a6e0 · 0b0a6e0
1 parent b0ade67
commit 0b0a6e0
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/brainforge/learner/abstract_learner.py b/brainforge/learner/abstract_learner.py
@@ -95,11 +95,13 @@ def evaluate_batch(self, x, y, metrics=()):
     def evaluate_stream(self, stream, steps, metrics=(), verbose=False):
         history = logging.MetricLogs.from_metric_list(steps, ["cost"], metrics)
         metrics = [_metrics.get(metric) for metric in metrics]
-        for x, y in stream:
+        for i, (x, y) in enumerate(stream, start=1):
             eval_metrics = self.evaluate_batch(x, y, metrics)
             history.record(eval_metrics)
             if verbose:
                 history.log("\r", end="")
+            if i >= steps:
+                break
         if verbose:
             print()
         history.reduce_mean()