🩹 Fix FFJTransform without context

probabilists · Feb 9, 2023 · 06089ed · 06089ed
1 parent 0ad5a1d
commit 06089ed
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/zuko/flows.py b/zuko/flows.py
@@ -869,7 +869,7 @@ def forward(self, y: Tensor = None) -> Transform:
         return FreeFormJacobianTransform(
             f=partial(self.f, y=y),
             time=self.time,
-            phi=(y, *self.ode.parameters()),
+            phi=self.ode.parameters() if y is None else (y, *self.ode.parameters()),
             exact=self.exact,
         )