Undo Return single acceleration for dyn mask

NKI-AI · Feb 1, 2024 · 55582eb · 55582eb
1 parent 94bfde4
commit 55582eb
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/direct/nn/mri_models.py b/direct/nn/mri_models.py
@@ -741,14 +741,22 @@ def perform_sampling(self, data: Dict[str, Any]) -> Dict[str, Any]:
             if "kspace" not in data:
                 raise ValueError("Expected data to contain key `kspace`, but not found.")
             sampling_model_kwargs = {"kspace": data["kspace"], "mask": data["sampling_mask"]}
+
+            acceleration = data["acceleration"]
+            center_fraction = data["center_fraction"]
+
+            if acceleration.dim > 1:
+                acceleration = acceleration[:, 0]
+                center_fraction = center_fraction[:, 0]
+
             sampling_model_kwargs.update(
                 filter_arguments_by_signature(
                     self.models["sampling_model"].forward,
                     {
                         "masked_kspace": data["masked_kspace"],
                         "sensitivity_map": data["sensitivity_map"],
-                        "acceleration": data["acceleration"],
-                        "center_fraction": data["center_fraction"],
+                        "acceleration": acceleration,
+                        "center_fraction": center_fraction,
                     },
                 )
             )