their latent bottleneck is tanh it seems, constraining it to -1 to 1 for flow matching in dynamics model. please open an issue if mistakened

dreamer4/dreamer4.py

@@ -291,16 +291,29 @@ class SwiGLUFeedforward(Module):
 
 class VideoTokenizer(Module):
     def __init__(
-        self
+        self,
+        dim,
+        dim_latent
     ):
         super().__init__()
 
+        self.encoded_to_latents = Sequential(
+            LinearNoBias(dim, dim_latent),
+            nn.Tanh(),
+        )
+
+        self.latents_to_decoder = LinearNoBias(dim_latent, dim)
+
+# dynamics model
+
 class DynamicsModel(Module):
     def __init__(
         self
     ):
         super().__init__()
 
+# dreamer
+
 class Dreamer(Module):
     def __init__(
         self,