loss_and_log fails when there is no loss

[jiyuuchc]

Ran into this bug in a rare edge case

in loss_and_logs.py:

    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) <--- should be jnp.array(0., float)
            losses_logs = {}