RandomEulerAngleDataset

Genentech · May 22, 2024 · 6395437 · 6395437
1 parent c17129f
commit 6395437
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 14 deletions.
diff --git a/src/beignet/datasets/__random_rotation_dataset.py b/src/beignet/datasets/__random_rotation_dataset.py
@@ -7,15 +7,16 @@
 
 
 class RandomRotationDataset(Dataset):
-    data: Tensor
-
     def __init__(
         self,
+        data: Tensor,
         *,
         transform: Callable | Transform | None = None,
     ) -> None:
         super().__init__()
 
+        self.data = data
+
         self.transform = transform
 
     def __getitem__(self, index: int) -> Tensor:

diff --git a/src/beignet/datasets/_random_euler_angle_dataset.py b/src/beignet/datasets/_random_euler_angle_dataset.py
@@ -23,16 +23,17 @@ def __init__(
         requires_grad: bool | None = False,
         transform: Callable | Transform | None = None,
     ) -> None:
-        self.data = beignet.random_euler_angle(
-            size,
-            axes,
-            degrees,
-            generator=generator,
-            dtype=dtype,
-            layout=layout,
-            device=device,
-            requires_grad=requires_grad,
-            pin_memory=pin_memory,
+        super().__init__(
+            beignet.random_euler_angle(
+                size,
+                axes,
+                degrees,
+                generator=generator,
+                dtype=dtype,
+                layout=layout,
+                device=device,
+                requires_grad=requires_grad,
+                pin_memory=pin_memory,
+            ),
+            transform=transform,
         )
-
-        super().__init__(transform=transform)