update mmnist

hrdae/dataloaders/datasets/moving_mnist.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 
-from torch import Tensor, cat, gather, int64, tensor
+from torch import Tensor, cat, gather, int64, tensor, device
+from torch.cuda import is_available
 from torch.utils.data import Dataset
 from torchvision import datasets
 
@@ -30,7 +31,11 @@ def __init__(self, *args, **kwargs) -> None:
 
     def __getitem__(self, idx: int) -> dict[str, Tensor]:
         # (n, h, w)
-        x_2d = super().__getitem__(idx).squeeze(1)
+        x_2d = self.data[idx].to(device("cuda:0") if is_available() else device("cpu"))
+        if self.transform is not None:
+            x_2d = self.transform(x_2d)
+        x_2d = x_2d.squeeze(1)
+
         n, h, w = x_2d.size()
 
         # (s,)

hrdae/dataloaders/transforms/normalization.py

@@ -1,7 +1,6 @@
 from dataclasses import dataclass
 
-from torch import Tensor, device
-from torch.cuda import is_available
+from torch import Tensor
 
 from .option import TransformOption
 
@@ -13,6 +12,4 @@ class MinMaxNormalizationOption(TransformOption):
 
 class MinMaxNormalization:
     def __call__(self, x: Tensor) -> Tensor:
-        x = x.to(device("cuda:0") if is_available() else device("cpu"))
-        x = (x - x.min()) / (x.max() - x.min()).cpu()
-        return x
+        return (x - x.min()) / (x.max() - x.min())

hrdae/models/networks/__init__.py

@@ -12,9 +12,7 @@
 from .r_dae import RDAE2dOption, RDAE3dOption, create_rdae2d, create_rdae3d
 
 
-def create_network(
-    out_channels: int, opt: NetworkOption
-) -> nn.Module:
+def create_network(out_channels: int, opt: NetworkOption) -> nn.Module:
     if (
         isinstance(opt, AutoEncoder2dNetworkOption)
         and type(opt) is AutoEncoder2dNetworkOption