From c7585c7bc7c2b53923944a9bb6af7a3664de478f Mon Sep 17 00:00:00 2001
From: James Fulton <djamesfulton@yahoo.co.uk>
Date: Tue, 21 Nov 2023 15:57:56 +0000
Subject: [PATCH] import fixes

---
 pvnet/data/datamodule.py           | 2 +-
 scripts/save_batches.py            | 4 +++-
 scripts/save_concurrent_batches.py | 4 +++-
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/pvnet/data/datamodule.py b/pvnet/data/datamodule.py
index 83f10487..ef38c791 100644
--- a/pvnet/data/datamodule.py
+++ b/pvnet/data/datamodule.py
@@ -5,6 +5,7 @@
 import torch
 from torch.utils.data import DataLoader
 from torch.utils.data.datapipes.datapipe import IterDataPipe
+from torch.utils.data.datapipes.iter import FileLister
 from torch.utils.data.datapipes._decorator import functional_datapipe
 
 from lightning.pytorch import LightningDataModule
@@ -184,7 +185,6 @@ def train_dataloader(self):
             datapipe = self._get_premade_batches_datapipe("train", shuffle=True)
         else:
             datapipe = self._get_datapipe(*self.train_period)
-        rs = MultiProcessingReadingService(**self.readingservice_config)
         return DataLoader(datapipe, **self._common_dataloader_kwargs)
 
     def val_dataloader(self):
diff --git a/scripts/save_batches.py b/scripts/save_batches.py
index a3526a76..c1884c02 100644
--- a/scripts/save_batches.py
+++ b/scripts/save_batches.py
@@ -25,11 +25,13 @@
 
 import hydra
 import torch
+from torch.utils.data import DataLoader
+from torch.utils.data.datapipes.datapipe import IterDataPipe
 from ocf_datapipes.training.pvnet import pvnet_datapipe
 from ocf_datapipes.utils.utils import stack_np_examples_into_batch
 from omegaconf import DictConfig, OmegaConf
 from sqlalchemy import exc as sa_exc
-from torch.utils.data import DataLoader
+
 from tqdm import tqdm
 
 from pvnet.data.datamodule import batch_to_tensor
diff --git a/scripts/save_concurrent_batches.py b/scripts/save_concurrent_batches.py
index 7fe22a38..e04a277c 100644
--- a/scripts/save_concurrent_batches.py
+++ b/scripts/save_concurrent_batches.py
@@ -23,6 +23,8 @@
 import hydra
 import numpy as np
 import torch
+from torch.utils.data import DataLoader
+from torch.utils.data.datapipes.datapipe import IterDataPipe
 from ocf_datapipes.training.common import (
     open_and_return_datapipes,
 )
@@ -31,7 +33,7 @@
 from ocf_datapipes.utils.utils import stack_np_examples_into_batch
 from omegaconf import DictConfig, OmegaConf
 from sqlalchemy import exc as sa_exc
-from torch.utils.data import DataLoader
+
 from tqdm import tqdm
 
 from pvnet.data.datamodule import batch_to_tensor