Small WandB improvements

Author: C-Achard

napari_cellseg3d/code_models/worker_training.py

@@ -787,7 +787,11 @@ def train(
             criterionW = None
             del criterionW
             torch.cuda.empty_cache()
+            if WANDB_INSTALLED:
+                wandb.finish()
         except Exception as e:
+            if WANDB_INSTALLED:
+                wandb.finish()
             msg = f"Training failed with exception: {e}"
             self.log(msg)
             self.raise_error(e, msg)
@@ -1726,6 +1730,9 @@ def get_patch_loader_func(num_samples):
             )
             self.log("Saving complete, exiting")
             model.to("cpu")
+
+            if WANDB_INSTALLED:
+                wandb.finish()
             # clear (V)RAM
             model = None
             del model
@@ -1741,6 +1748,8 @@ def get_patch_loader_func(num_samples):
                 torch.cuda.empty_cache()
 
         except Exception as e:
+            if WANDB_INSTALLED:
+                wandb.finish()
             self.raise_error(e, "Error in training")
             self.quit()
         finally: