format tests, remove print

microsoft · Sep 24, 2024 · fde215c · fde215c
1 parent 034c647
commit fde215c
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 19 deletions.
diff --git a/tests/samplers/test_batch.py b/tests/samplers/test_batch.py
@@ -148,12 +148,10 @@ def test_weighted_sampling(self) -> None:
     def test_random_seed(self) -> None:
         ds = CustomGeoDataset()
         ds.index.insert(0, (0, 10, 0, 10, 0, 10))
-        sampler1 = RandomBatchGeoSampler(
-            ds, 1, 1, generator=torch.Generator().manual_seed(0)
-        )
-        sampler2 = RandomBatchGeoSampler(
-            ds, 1, 1, generator=torch.Generator().manual_seed(0)
-        )
+        generator1 = torch.Generator().manual_seed(0)
+        generator2 = torch.Generator().manual_seed(0)
+        sampler1 = RandomBatchGeoSampler(ds, 1, 1, generator=generator1)
+        sampler2 = RandomBatchGeoSampler(ds, 1, 1, generator=generator2)
         sample1 = next(iter(sampler1))
         sample2 = next(iter(sampler2))
         assert sample1 == sample2

diff --git a/tests/samplers/test_single.py b/tests/samplers/test_single.py
@@ -143,12 +143,10 @@ def test_weighted_sampling(self) -> None:
     def test_random_seed(self) -> None:
         ds = CustomGeoDataset()
         ds.index.insert(0, (0, 10, 0, 10, 0, 10))
-        sampler1 = RandomGeoSampler(
-            ds, 1, 1, generator=torch.Generator().manual_seed(0)
-        )
-        sampler2 = RandomGeoSampler(
-            ds, 1, 1, generator=torch.Generator().manual_seed(0)
-        )
+        generator1 = torch.Generator().manual_seed(0)
+        generator2 = torch.Generator().manual_seed(0)
+        sampler1 = RandomGeoSampler(ds, 1, 1, generator=generator1)
+        sampler2 = RandomGeoSampler(ds, 1, 1, generator=generator2)
         sample1 = next(iter(sampler1))
         sample2 = next(iter(sampler2))
         assert sample1 == sample2
@@ -306,13 +304,11 @@ def test_shuffle_seed(self) -> None:
         ds = CustomGeoDataset()
         ds.index.insert(0, (0, 10, 0, 10, 0, 10))
         ds.index.insert(1, (0, 11, 0, 11, 0, 11))
-        sampler1 = PreChippedGeoSampler(
-            ds, shuffle=True, generator=torch.Generator().manual_seed(0)
-        )
+        generator1 = torch.Generator().manual_seed(0)
+        generator2 = torch.Generator().manual_seed(0)
+        sampler1 = PreChippedGeoSampler(ds, shuffle=True, generator=generator1)
+        sampler2 = PreChippedGeoSampler(ds, shuffle=True, generator=generator2)
         sample1 = next(iter(sampler1))
-        sampler2 = PreChippedGeoSampler(
-            ds, shuffle=True, generator=torch.Generator().manual_seed(0)
-        )
         sample2 = next(iter(sampler2))
         assert sample1 == sample2
 

diff --git a/torchgeo/samplers/single.py b/torchgeo/samplers/single.py
@@ -320,7 +320,6 @@ def __iter__(self) -> Iterator[BoundingBox]:
             generator = partial(torch.randperm, generator=self.generator)
 
         for idx in generator(len(self)):
-            print(idx)
             yield BoundingBox(*self.hits[idx].bounds)
 
     def __len__(self) -> int: