Fix list-based indexing in the movement module.

PiperOrigin-RevId: 568791855
google · Sep 27, 2023 · 0b30abb · 0b30abb
1 parent 6c503de
commit 0b30abb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ffn/inference/movement.py b/ffn/inference/movement.py
@@ -71,7 +71,7 @@ def get_scored_move_offsets(deltas, prob_map, threshold=0.9):
       # of the subvolume orthogonal to the current axis.
       face_sel = subvol_sel[:]
       face_sel[axis] = axis_offset + center[axis]
-      face_prob = prob_map[face_sel]
+      face_prob = prob_map[tuple(face_sel)]
       shape = face_prob.shape
 
       # Find voxel with maximum activation.