if logging with wandb, unwatch the model after training

sorenmacbeth · sorenmacbeth · commit a7aa3daa738b · 2024-12-02T14:07:32.000-08:00
diff --git a/src/pytorch_tabular/tabular_model.py b/src/pytorch_tabular/tabular_model.py
@@ -686,6 +686,8 @@ def train(
                 "/n" + "Original Error: " + oom_handler.oom_msg
             )
         self._is_fitted = True
+        if self.track_experiment and self.config.log_target == "wandb":
+            self.logger.experiment.unwatch(self.model)
         if self.verbose:
             logger.info("Training the model completed")
         if self.config.load_best:

Original file line number	Diff line number	Diff line change
`@@ -686,6 +686,8 @@ def train(`
`686`	`686`	`"/n" + "Original Error: " + oom_handler.oom_msg`
`687`	`687`	`)`
`688`	`688`	`self._is_fitted = True`
	`689`	`+ if self.track_experiment and self.config.log_target == "wandb":`
	`690`	`+ self.logger.experiment.unwatch(self.model)`
`689`	`691`	`if self.verbose:`
`690`	`692`	`logger.info("Training the model completed")`
`691`	`693`	`if self.config.load_best:`