From d9bc1da907915fabe1410b4524e9642edf5a3600 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andreas=20S=C3=B8gaard?= <andreas.sogaard@gmail.com>
Date: Wed, 31 Aug 2022 13:28:26 +0200
Subject: [PATCH 1/2] Create wandb directory if it doesn't exist

---
 examples/train_model.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/examples/train_model.py b/examples/train_model.py
index 67779e302..49439f8ba 100644
--- a/examples/train_model.py
+++ b/examples/train_model.py
@@ -1,4 +1,4 @@
-import os.path
+import os
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import EarlyStopping
@@ -41,6 +41,9 @@
     log_model=True,
 )
 
+# -- Ensure that custom output directory exists
+os.makedirs(wandb_logger.save_dir, exists_ok=True)
+
 
 # Main function definition
 def main():

From 088632e765bd6ef71581eb817dae7455c6b155c7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andreas=20S=C3=B8gaard?= <andreas.sogaard@gmail.com>
Date: Wed, 31 Aug 2022 13:32:22 +0200
Subject: [PATCH 2/2] Reorder

---
 examples/train_model.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/examples/train_model.py b/examples/train_model.py
index 49439f8ba..f2a0f417e 100644
--- a/examples/train_model.py
+++ b/examples/train_model.py
@@ -33,17 +33,18 @@
 features = FEATURES.DEEPCORE
 truth = TRUTH.DEEPCORE[:-1]
 
+# Make sure W&B output directory exists
+WANDB_DIR = "./wandb/"
+os.makedirs(WANDB_DIR, exist_ok=True)
+
 # Initialise Weights & Biases (W&B) run
 wandb_logger = WandbLogger(
     project="example-script",
     entity="graphnet-team",
-    save_dir="./wandb/",
+    save_dir=WANDB_DIR,
     log_model=True,
 )
 
-# -- Ensure that custom output directory exists
-os.makedirs(wandb_logger.save_dir, exists_ok=True)
-
 
 # Main function definition
 def main():