Return registered image if registration

NKI-AI · Aug 13, 2024 · cbf5a56 · cbf5a56
1 parent 8742423
commit cbf5a56
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/direct/nn/vsharp/vsharp_engine.py b/direct/nn/vsharp/vsharp_engine.py
@@ -161,7 +161,7 @@ def _do_iteration(
 
         output_image = output_images[-1]
         return DoIterationOutput(
-            output_image=output_image,
+            output_image=(output_image, registered_image) if "registration_model" in self.models else output_image,
             sensitivity_map=data["sensitivity_map"],
             sampling_mask=data["sampling_mask"],
             data_dict={**loss_dict},