diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py
index dca0796..d16e12d 100644
--- a/dreamer4/dreamer4.py
+++ b/dreamer4/dreamer4.py
@@ -1883,6 +1883,8 @@ class DynamicsWorldModel(Module):
         policy_optim: Optimizer | None = None,
         value_optim: Optimizer | None = None
     ):
+        assert experience.is_batched
+
         latents = experience.latents
         actions = experience.actions
         old_log_probs = experience.log_probs
diff --git a/pyproject.toml b/pyproject.toml
index 094eab5..e65affc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "dreamer4"
-version = "0.0.45"
+version = "0.0.46"
 description = "Dreamer 4"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }