Remove device option from TorchDiskDataset

src/tinytopics/data.py

@@ -114,13 +114,11 @@ def __init__(
         self,
         data: str | Path,
         indices: Sequence[int] | None = None,
-        device: torch.device | None = None,
     ) -> None:
         """
         Args:
             data: Path to `.pt` file (str or Path).
             indices: Optional sequence of indices to use as valid indices.
-            device: Optional device to load tensors to. Defaults to CPU.
         """
         self.data_path = Path(data)
         if not self.data_path.exists():
@@ -147,7 +145,6 @@ def __init__(
             self.mmap_supported = False
 
         self.indices = indices or range(self.shape[0])
-        self.device = device or torch.device("cpu")
 
     def __len__(self) -> int:
         return len(self.indices)
@@ -160,9 +157,9 @@ def __getitem__(self, idx: int) -> torch.Tensor:
                 self.mmap_data = torch.load(
                     self.data_path, map_location="cpu", weights_only=True, mmap=True
                 )
-            return self.mmap_data[real_idx].to(self.device)
+            return self.mmap_data[real_idx]
         else:
-            return self.data[real_idx].to(self.device)
+            return self.data[real_idx]
 
     @property
     def num_terms(self) -> int:

tests/test_data.py

@@ -109,7 +109,6 @@ def test_torch_disk_dataset_from_file(tmp_path):
         assert isinstance(item, torch.Tensor)
         assert item.shape == (5,)
         assert torch.allclose(item, data[i])
-        assert item.device == torch.device("cpu")
 
 
 def test_torch_disk_dataset_with_indices(tmp_path):
@@ -134,39 +133,6 @@ def test_torch_disk_dataset_with_indices(tmp_path):
         assert torch.allclose(item, data[orig_idx])
 
 
-def test_torch_disk_dataset_with_device(tmp_path):
-    """Test TorchDiskDataset with specific device."""
-    data = torch.rand(10, 5, dtype=torch.float32)
-    file_path = tmp_path / "test_data.pt"
-    torch.save(data, file_path)
-
-    device = torch.device("cpu")  # Use CPU for testing
-    dataset = TorchDiskDataset(file_path, device=device)
-
-    # Test data access with device
-    for i in range(len(dataset)):
-        item = dataset[i]
-        assert item.device == device
-        assert torch.allclose(item, data[i])
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA device not available")
-def test_torch_disk_dataset_cuda_device(tmp_path):
-    """Test TorchDiskDataset with CUDA device."""
-    data = torch.rand(10, 5, dtype=torch.float32)
-    file_path = tmp_path / "test_data.pt"
-    torch.save(data, file_path)
-
-    device = torch.device("cuda")
-    dataset = TorchDiskDataset(file_path, device=device)
-
-    # Test data access with CUDA device
-    for i in range(len(dataset)):
-        item = dataset[i]
-        assert item.device == device
-        assert torch.allclose(item.cpu(), data[i])
-
-
 def test_torch_disk_dataset_file_not_found():
     """Test TorchDiskDataset with non-existent file."""
     with pytest.raises(FileNotFoundError):