update search with int flow

gauenk · Oct 26, 2023 · 8b9e4c4 · 8b9e4c4
1 parent c434272
commit 8b9e4c4
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/lib/stnls/search/non_local_search.py b/lib/stnls/search/non_local_search.py
@@ -57,7 +57,10 @@ def nls_forward(vid0, vid1, flows,
 
     # -- forward --
     if itype == "int":
-        flows = flows.round().int()
+        if flows.dtype != th.int:
+            flows = flows.round().int()
+        else:
+            flows = flows.int()
         inds = inds.int()
         stride1 = max(1,int(stride1))
         fwd_fxn = stnls_cuda.non_local_search_int_forward