arg type

microsoft · Feb 18, 2024 · 25c8d64 · 25c8d64
1 parent 2e0c93e
commit 25c8d64
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/torchgeo/datamodules/cdlsentinel2.py b/torchgeo/datamodules/cdlsentinel2.py
@@ -87,7 +87,9 @@ def setup(self, stage: str) -> None:
         generator = torch.Generator().manual_seed(0)
 
         (self.train_dataset, self.val_dataset, self.test_dataset) = (
-            random_grid_cell_assignment(self.dataset, [0.8, 0.1, 0.1], generator)
+            random_grid_cell_assignment(
+                self.dataset, [0.8, 0.1, 0.1], generator=generator
+            )
         )
 
         if stage in ["fit"]: