Bugfix concatenation axis

bayesflow-org · Jun 12, 2024 · 42d20ab · 42d20ab
1 parent 047204d
commit 42d20ab
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bayesflow/experimental/utils/dictutils.py b/bayesflow/experimental/utils/dictutils.py
@@ -19,7 +19,7 @@ def concatenate_tensors(tensor_dict: dict[str, Tensor], filter_list: list, axis:
     An optional axis can be specified (default: last axis).
     """
 
-    return ops.concatenate([v for k, v in tensor_dict.items() if k in filter_list])
+    return ops.concatenate([v for k, v in tensor_dict.items() if k in filter_list], axis=axis)
 
 
 def keras_kwargs(kwargs: dict):