1 files changed, 7 insertions, 3 deletions

diff --git a/data/csv.py b/data/csv.py
index af36d9e..e901ab4 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -59,7 +59,7 @@ class CSVDataModule():
         center_crop: bool = False,
         template_key: str = "template",
         valid_set_size: Optional[int] = None,
-        generator: Optional[torch.Generator] = None,
+        seed: Optional[int] = None,
         filter: Optional[Callable[[CSVDataItem], bool]] = None,
         collate_fn=None,
         num_workers: int = 0
@@ -84,7 +84,7 @@ class CSVDataModule():
         self.template_key = template_key
         self.interpolation = interpolation
         self.valid_set_size = valid_set_size
-        self.generator = generator
+        self.seed = seed
         self.filter = filter
         self.collate_fn = collate_fn
         self.num_workers = num_workers
@@ -155,7 +155,11 @@ class CSVDataModule():
         valid_set_size = max(valid_set_size, 1)
         train_set_size = num_images - valid_set_size
 
-        data_train, data_val = random_split(items, [train_set_size, valid_set_size], self.generator)
+        generator = torch.Generator(device="cpu")
+        if self.seed is not None:
+            generator = generator.manual_seed(self.seed)
+
+        data_train, data_val = random_split(items, [train_set_size, valid_set_size], generator=generator)
 
         self.data_train = self.pad_items(data_train, self.num_class_images)
         self.data_val = self.pad_items(data_val)