FlowMatching: build self instead of base_distribution

bayesflow-org · Feb 11, 2025 · 1081bbd · 1081bbd
1 parent fdac500
commit 1081bbd
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/bayesflow/networks/flow_matching/flow_matching.py b/bayesflow/networks/flow_matching/flow_matching.py
@@ -191,9 +191,10 @@ def compute_metrics(
         else:
             # not pre-configured, resample
             x1 = x
-            if not self.base_distribution.built:
-                # ensure that base distribution is built
-                self.base_distribution.build(keras.ops.shape(x1))
+            if not self.built:
+                xz_shape = keras.ops.shape(x1)
+                conditions_shape = None if conditions is None else keras.ops.shape(conditions)
+                self.build(xz_shape, conditions_shape)
             x0 = self.base_distribution.sample(keras.ops.shape(x1)[:-1])
 
             if self.use_optimal_transport: