diff --git a/Ovi/vae/vae.py b/Ovi/vae/vae.py
index d1b9254..0e0d207 100644
--- a/Ovi/vae/vae.py
+++ b/Ovi/vae/vae.py
@@ -75,14 +75,21 @@ class VAE(nn.Module):
         super().__init__()
 
         if data_dim == 80:
-            self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_80D, dtype=torch.float32))
-            self.data_std = nn.Buffer(torch.tensor(DATA_STD_80D, dtype=torch.float32))
+            data_mean = torch.tensor(DATA_MEAN_80D, dtype=torch.float32)
+            data_std  = torch.tensor(DATA_STD_80D, dtype=torch.float32)
         elif data_dim == 128:
-            self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_128D, dtype=torch.float32))
-            self.data_std = nn.Buffer(torch.tensor(DATA_STD_128D, dtype=torch.float32))
+            data_mean = torch.tensor(DATA_MEAN_128D, dtype=torch.float32)
+            data_std  = torch.tensor(DATA_STD_128D, dtype=torch.float32)
+        else:
+            raise ValueError(f"Unsupported data_dim={data_dim}, expected 80 or 128")
 
-        self.data_mean = self.data_mean.view(1, -1, 1)
-        self.data_std = self.data_std.view(1, -1, 1)
+        # match old shape: (1, channels, 1)
+        data_mean = data_mean.view(1, -1, 1)
+        data_std  = data_std.view(1, -1, 1)
+
+        # register as buffers so they move with .to(device) / .cuda()
+        self.register_buffer("data_mean", data_mean)
+        self.register_buffer("data_std", data_std)
 
         self.encoder = Encoder1D(
             dim=hidden_dim,