fix input shape

Author: arrjon

bayesflow/networks/consistency_models/consistency_model.py

@@ -168,11 +168,12 @@ def build(self, xz_shape, conditions_shape=None):
 
         input_shape = list(xz_shape)
 
-        # time vector
-        input_shape[-1] += 1
+        if self._concatenate_subnet_input:
+            # time vector
+            input_shape[-1] += 1
 
-        if conditions_shape is not None:
-            input_shape[-1] += conditions_shape[-1]
+            if conditions_shape is not None:
+                input_shape[-1] += conditions_shape[-1]
 
         input_shape = tuple(input_shape)
 

bayesflow/networks/diffusion_model/diffusion_model.py

@@ -135,10 +135,11 @@ def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
         self.output_projector.units = xz_shape[-1]
         input_shape = list(xz_shape)
 
-        # construct time vector
-        input_shape[-1] += 1
-        if conditions_shape is not None:
-            input_shape[-1] += conditions_shape[-1]
+        if self._concatenate_subnet_input:
+            # construct time vector
+            input_shape[-1] += 1
+            if conditions_shape is not None:
+                input_shape[-1] += conditions_shape[-1]
 
         input_shape = tuple(input_shape)
 

bayesflow/networks/flow_matching/flow_matching.py

@@ -128,11 +128,12 @@ def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
 
         self.output_projector.units = xz_shape[-1]
 
-        # account for concatenating the time and conditions
         input_shape = list(xz_shape)
-        input_shape[-1] += 1
-        if conditions_shape is not None:
-            input_shape[-1] += conditions_shape[-1]
+        if self._concatenate_subnet_input:
+            # account for concatenating the time and conditions
+            input_shape[-1] += 1
+            if conditions_shape is not None:
+                input_shape[-1] += conditions_shape[-1]
         input_shape = tuple(input_shape)
 
         self.subnet.build(input_shape)