fix wandb solutions plotting

SkafteNicki · Jan 14, 2025 · 868badc · 868badc
1 parent e1da4cf
commit 868badc
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/s4_debugging_and_logging/exercise_files/weights_and_bias_solution2.py b/s4_debugging_and_logging/exercise_files/weights_and_bias_solution2.py
@@ -68,8 +68,9 @@ def train(lr: float = 0.001, batch_size: int = 32, epochs: int = 5) -> None:
                 plot_chance_level=(class_id == 2),
             )
 
+        # alternatively use wandb.log({"roc": wandb.Image(plt)}
         wandb.plot({"roc": plt})
-        # alternatively the wandb.plot.roc_curve function can be used
+        plt.close()  # close the plot to avoid memory leaks and overlapping figures
 
 
 if __name__ == "__main__":

diff --git a/s4_debugging_and_logging/exercise_files/weights_and_bias_solution3.py b/s4_debugging_and_logging/exercise_files/weights_and_bias_solution3.py
@@ -68,8 +68,9 @@ def train(lr: float = 0.001, batch_size: int = 32, epochs: int = 5) -> None:
                 plot_chance_level=(class_id == 2),
             )
 
+        # alternatively use wandb.log({"roc": wandb.Image(plt)}
         wandb.plot({"roc": plt})
-        # alternatively the wandb.plot.roc_curve function can be used
+        plt.close()  # close the plot to avoid memory leaks and overlapping figures
 
     final_accuracy = accuracy_score(targets, preds.argmax(dim=1))
     final_precision = precision_score(targets, preds.argmax(dim=1), average="weighted")