From 10835e9aaa577569be1a427888750c2966a3f75c Mon Sep 17 00:00:00 2001 From: joeloskarsson Date: Fri, 7 Jun 2024 16:02:42 +0200 Subject: [PATCH] Do not include directory structure of data config in wandb storage --- neural_lam/models/ar_model.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py index 1dd6b72c..b8f0bfd7 100644 --- a/neural_lam/models/ar_model.py +++ b/neural_lam/models/ar_model.py @@ -215,8 +215,7 @@ def training_step(self, batch): def on_train_start(self): """Save data config file to wandb at start of training""" - if self.trainer.is_global_zero: - wandb.save(self.args.data_config) + self.save_data_config() def all_gather_cat(self, tensor_to_gather): """ @@ -528,8 +527,7 @@ def aggregate_and_plot_metrics(self, metrics_dict, prefix): def on_test_start(self): """Save data config file to wandb at start of test""" - if self.trainer.is_global_zero: - wandb.save(self.args.data_config) + self.save_data_config() def on_test_epoch_end(self): """ @@ -607,3 +605,12 @@ def on_load_checkpoint(self, checkpoint): if not self.restore_opt: opt = self.configure_optimizers() checkpoint["optimizer_states"] = [opt.state_dict()] + + def save_data_config(self): + """Save data config file to wandb""" + if self.trainer.is_global_zero: + wandb.save( + self.args.data_config, + base_path=os.path.dirname(self.args.data_config), + policy="now", + )