Do not include directory structure of data config in wandb storage

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",
+            )