fix batch_shape in sample

Author: arrjon

bayesflow/approximators/continuous_approximator.py

@@ -535,10 +535,9 @@ def _sample(
             inference_conditions = keras.ops.broadcast_to(
                 inference_conditions, (batch_size, num_samples, *keras.ops.shape(inference_conditions)[2:])
             )
-            batch_shape = (
-                batch_size,
-                num_samples,
-            )
+
+            target_dim = self.inference_network.base_distribution.dims
+            batch_shape = keras.ops.shape(inference_conditions)[: -len(target_dim)]
         else:
             batch_shape = (num_samples,)