diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py
index bf1dc59..e3d16da 100644
--- a/dreamer4/dreamer4.py
+++ b/dreamer4/dreamer4.py
@@ -2543,6 +2543,7 @@ class DynamicsWorldModel(Module):
         batch_size = 1,
         agent_index = 0,
         tasks: int | Tensor | None = None,
+        latent_gene_ids = None,
         image_height = None,
         image_width = None,
         return_decoded_video = None,
@@ -2658,6 +2659,7 @@ class DynamicsWorldModel(Module):
                     step_sizes = step_size,
                     rewards = decoded_rewards,
                     tasks = tasks,
+                    latent_gene_ids = latent_gene_ids,
                     discrete_actions = decoded_discrete_actions,
                     continuous_actions = decoded_continuous_actions,
                     proprio = noised_proprio_with_context,
diff --git a/pyproject.toml b/pyproject.toml
index 5034a63..a8d2968 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "dreamer4"
-version = "0.0.89"
+version = "0.0.90"
 description = "Dreamer 4"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }