turn python values into scalar tensors

princeton-vl · Oct 11, 2021 · c78cf64 · c78cf64
1 parent a09fffc
commit c78cf64
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/core/raft.py b/core/raft.py
@@ -83,7 +83,7 @@ def upsample_flow(self, flow, mask):
         return up_flow.reshape(N, 2, 8*H, 8*W)
 
 
-    def forward(self, image1, image2, iters=12, flow_init=None, upsample=True, test_mode=False):
+    def forward(self, image1, image2, iters=torch.tensor(12), flow_init=torch.tensor([]), upsample=torch.tensor(True), test_mode=torch.tensor(False)):
         """ Estimate optical flow between pair of frames """
 
         image1 = 2 * (image1 / 255.0) - 1.0
@@ -115,7 +115,7 @@ def forward(self, image1, image2, iters=12, flow_init=None, upsample=True, test_
 
         coords0, coords1 = self.initialize_flow(image1)
 
-        if flow_init is not None:
+        if flow_init is not None and flow_init.numel()>0:
             coords1 = coords1 + flow_init
 
         flow_predictions = []

diff --git a/demo.py b/demo.py
@@ -61,7 +61,7 @@ def demo(args):
             padder = InputPadder(image1.shape)
             image1, image2 = padder.pad(image1, image2)
 
-            flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
+            flow_low, flow_up = model(image1, image2, iters=torch.tensor(20), test_mode=torch.tensor(True))
             viz(image1, flow_up)