Refactor BoringFabric in tests (Lightning-AI#19364)

tests/tests_fabric/helpers/datasets.py

@@ -0,0 +1,27 @@
+from typing import Iterator
+
+import torch
+from torch import Tensor
+from torch.utils.data import Dataset, IterableDataset
+
+
+class RandomDataset(Dataset):
+    def __init__(self, size: int, length: int) -> None:
+        self.len = length
+        self.data = torch.randn(length, size)
+
+    def __getitem__(self, index: int) -> Tensor:
+        return self.data[index]
+
+    def __len__(self) -> int:
+        return self.len
+
+
+class RandomIterableDataset(IterableDataset):
+    def __init__(self, size: int, count: int) -> None:
+        self.count = count
+        self.size = size
+
+    def __iter__(self) -> Iterator[Tensor]:
+        for _ in range(self.count):
+            yield torch.randn(self.size)

tests/tests_fabric/strategies/test_deepspeed_integration.py

@@ -25,7 +25,7 @@
 from lightning.fabric.strategies import DeepSpeedStrategy
 from torch.utils.data import DataLoader
 
-from tests_fabric.helpers.models import RandomDataset, RandomIterableDataset
+from tests_fabric.helpers.datasets import RandomDataset, RandomIterableDataset
 from tests_fabric.helpers.runif import RunIf
 from tests_fabric.test_fabric import BoringModel