Describe the bug
The location values of probability mask is not being used correctly in the numpy output. The locations are a numpy array (instead of a tuple) so when used in indexing an ndarray, they are not being interpreted as the location in two/three dimension. Instead they are treated as the index of the first dimension. Explicit conversion to tuple can solve this issue.