import fixes

openclimatefix · Nov 21, 2023 · 4e1cd7b · 4e1cd7b
1 parent d74d7ed
commit 4e1cd7b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/scripts/save_batches.py b/scripts/save_batches.py
@@ -30,6 +30,7 @@
 from omegaconf import DictConfig, OmegaConf
 from sqlalchemy import exc as sa_exc
 from torch.utils.data import DataLoader
+from torch.utils.data.datapipes.iter import IterableWrapper
 from tqdm import tqdm
 
 from pvnet.data.datamodule import batch_to_tensor
@@ -80,7 +81,7 @@ def _save_batches_with_dataloader(batch_pipe, batch_dir, num_batches, dataloader
 
 @hydra.main(config_path="../configs/", config_name="config.yaml", version_base="1.2")
 def main(config: DictConfig):
-    "Constructs and saves validation and training batches."
+    """Constructs and saves validation and training batches."""
     config_dm = config.datamodule
 
     print_config(config, resolve=False)

diff --git a/scripts/save_concurrent_batches.py b/scripts/save_concurrent_batches.py
@@ -32,6 +32,7 @@
 from omegaconf import DictConfig, OmegaConf
 from sqlalchemy import exc as sa_exc
 from torch.utils.data import DataLoader
+from torch.utils.data.datapipes.iter import IterableWrapper
 from tqdm import tqdm
 
 from pvnet.data.datamodule import batch_to_tensor