File tree 1 file changed +2
-2
lines changed
1 file changed +2
-2
lines changed Original file line number Diff line number Diff line change @@ -61,7 +61,7 @@ def sample_method( # pylint: disable=no-self-use
61
61
num_images: number of images to sample over
62
62
mask: mask of possible pixels in an image to sample from.
63
63
"""
64
- if mask :
64
+ if isinstance ( mask , torch . Tensor ) :
65
65
nonzero_indices = torch .nonzero (mask [..., 0 ], as_tuple = False )
66
66
chosen_indices = random .sample (range (len (nonzero_indices )), k = batch_size )
67
67
indices = nonzero_indices [chosen_indices ]
@@ -222,7 +222,7 @@ def sample_method( # pylint: disable=no-self-use
222
222
device : Union [torch .device , str ] = "cpu" ,
223
223
) -> TensorType ["batch_size" , 3 ]:
224
224
225
- if mask :
225
+ if isinstance ( mask , torch . Tensor ) :
226
226
# Note: if there is a mask, sampling reduces back to uniform sampling, which gives more
227
227
# sampling weight to the poles of the image than the equators.
228
228
# TODO(kevinddchen): implement the correct mask-sampling method.
You can’t perform that action at this time.
0 commit comments