fixing early return returning a float and not a tuple

netam/framework.py

@@ -775,8 +775,8 @@ def joint_train(
             )
             self.reset_optimization(new_lr)
             loss_history_l.append(self.train(epochs, out_prefix=out_prefix))
-            if cycle < cycle_count - 1:
-                optimize_branch_lengths()
+            # We standardize and optimize the branch lengths after each cycle, even the last one.
+            optimize_branch_lengths()
             self.mark_branch_lengths_optimized(cycle + 1)
 
         return pd.concat(loss_history_l, ignore_index=True)
@@ -932,7 +932,7 @@ def _find_optimal_branch_length(
         **optimization_kwargs,
     ):
         if torch.sum(mutation_indicator) == 0:
-            return 0.0
+            return 0.0, False
 
         rates, _ = self.model(
             encoded_parent.unsqueeze(0),