Fix LossAndLog error when there is no loss. (#75) (#79)

Thanks! Fix seems good. CI is failing, same stuff that happened on Treeo.
cgarciae · Feb 26, 2023 · a63a5a9 · a63a5a9
1 parent 551f768
commit a63a5a9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/treex/metrics/loss_and_logs.py b/treex/metrics/loss_and_logs.py
@@ -94,7 +94,7 @@ def compute(self) -> tp.Tuple[jnp.ndarray, Logs, Logs]:
         if self.losses is not None:
             loss, losses_logs = self.losses.compute()
         else:
-            loss = jnp.zeros(0.0, dtype=jnp.float32)
+            loss = jnp.array(0.0, dtype=jnp.float32)
             losses_logs = {}
 
         if self.metrics is not None: