oops

dreamer4/dreamer4.py

@@ -1251,6 +1251,7 @@ class AxialSpaceTimeTransformer(Module):
         final_norm = True
     ):
         super().__init__()
+        assert depth >= time_block_every, f'depth must be at least {time_block_every}'
 
         # hyper connections
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dreamer4"
-version = "0.0.55"
+version = "0.0.56"
 description = "Dreamer 4"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

tests/test_dreamer.py

@@ -483,7 +483,7 @@ def test_cache_generate():
         max_steps = 64,
         num_tasks = 4,
         num_latent_tokens = 4,
-        depth = 1,
+        depth = 4,
         num_spatial_tokens = 1,
         pred_orig_latent = True,
         num_discrete_actions = 4,