Make number of training samples per year configurable for FCN AFNO ex…

…ample (#115) * Update config.yaml * Update train_era5.py
NVIDIA · Sep 28, 2023 · 674086c · 674086c
1 parent 0dd6368
commit 674086c
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/examples/weather/fcn_afno/conf/config.yaml b/examples/weather/fcn_afno/conf/config.yaml
@@ -21,3 +21,5 @@ hydra:
 wb_artifacts: False
 
 start_epoch: 0
+
+num_samples_per_year_train: 1456
diff --git a/examples/weather/fcn_afno/train_era5.py b/examples/weather/fcn_afno/train_era5.py
@@ -116,7 +116,7 @@ def main(cfg: DictConfig) -> None:
         data_dir="/data/train/",
         stats_dir="/data/stats/",
         channels=[i for i in range(20)],
-        num_samples_per_year=1456,  # Need better shard fix
+        num_samples_per_year=cfg.num_samples_per_year_train,
         batch_size=2,
         patch_size=(8, 8),
         num_workers=8,
Original file line number	Diff line number	Diff line change
Expand Up		@@ -21,3 +21,5 @@ hydra:
		wb_artifacts: False

		start_epoch: 0

		num_samples_per_year_train: 1456