From 372b943b007a43fd573ba02b4b4b5ad1868907e8 Mon Sep 17 00:00:00 2001
From: Andreas Schuh <aschuh@heartflow.com>
Date: Tue, 6 Feb 2024 01:09:06 +0000
Subject: [PATCH] [core] Multiply input by 'scale' even when 'steps=0' in
 expv()

---
 src/deepali/core/flow.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/deepali/core/flow.py b/src/deepali/core/flow.py
index 0f89d4e..e9290b4 100644
--- a/src/deepali/core/flow.py
+++ b/src/deepali/core/flow.py
@@ -357,6 +357,8 @@ def expv(
     if steps < 0:
         raise ValueError("expv() 'steps' must be positive value")
     if steps == 0:
+        if abs(scale - 1) > 1e-15:
+            flow = flow.mul(scale)
         return flow
     device = flow.device
     grid = Grid(shape=flow.shape[2:], align_corners=align_corners)