guard for total_steps = 0

Author: LarsKue

bayesflow/approximators/backend_approximators/jax_approximator.py

@@ -76,6 +76,9 @@ def _aggregate_fn(_logs, _step_logs):
             return keras.tree.map_structure(keras.ops.add, _logs, _step_logs)
 
         def _reduce_fn(_logs, _total_steps):
+            if total_steps == 0:
+                return _logs
+
             def _div(val):
                 return val / _total_steps
 

bayesflow/approximators/backend_approximators/tensorflow_approximator.py

@@ -75,6 +75,9 @@ def _aggregate_fn(_logs, _step_logs):
             return keras.tree.map_structure(keras.ops.add, _logs, _step_logs)
 
         def _reduce_fn(_logs, _total_steps):
+            if total_steps == 0:
+                return _logs
+
             def _div(val):
                 return val / _total_steps
 

bayesflow/approximators/backend_approximators/torch_approximator.py

@@ -76,6 +76,9 @@ def _aggregate_fn(_logs, _step_logs):
             return keras.tree.map_structure(keras.ops.add, _logs, _step_logs)
 
         def _reduce_fn(_logs, _total_steps):
+            if total_steps == 0:
+                return _logs
+
             def _div(val):
                 return val / _total_steps