torchworld/structures/grid: added device and len methods

d4l3k · Oct 27, 2023 · 5d3a7da · 5d3a7da
1 parent b4665c8
commit 5d3a7da
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/torchworld/structures/grid.py b/torchworld/structures/grid.py
@@ -12,6 +12,9 @@
 
 @dataclass
 class BaseGrid(ABC):
+    data: torch.Tensor
+    time: torch.Tensor
+
     @abstractmethod
     def to(self: T, target: Union[torch.device, str]) -> T:
         ...
@@ -22,6 +25,13 @@ def cuda(self) -> T:
     def cpu(self) -> T:
         return self.to(torch.device("cpu"))
 
+    @property
+    def device(self) -> torch.device:
+        return self.data.device
+
+    def __len__(self) -> int:
+        return len(self.data)
+
 
 @dataclass
 class Grid3d(BaseGrid):

diff --git a/torchworld/structures/test_grid.py b/torchworld/structures/test_grid.py
@@ -18,6 +18,9 @@ def test_grid_3d(self) -> None:
         grid = grid.to("cpu")
         grid = grid.cpu()
 
+        self.assertEqual(len(grid), 2)
+        self.assertEqual(grid.device, grid.data.device)
+
     def test_grid_image(self) -> None:
         grid = GridImage(
             data=torch.rand(2, 3, 4, 5),