From c18c624be6f477276c7f78b4f50005b57ba4219c Mon Sep 17 00:00:00 2001
From: lucidrains <lucidrains@gmail.com>
Date: Wed, 1 Oct 2025 10:39:16 -0700
Subject: [PATCH] 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 | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py
index 07643b7..be2aa43 100644
--- a/dreamer4/dreamer4.py
+++ b/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,