Slightly modify the code to avoid warning. (#737)

`UserWarning: indexing with dtype torch.uint8 is now deprecated, please use a dtype torch.bool instead.`

Co-authored-by: Brian Johnson <[email protected]>
Co-authored-by: holly1238 <[email protected]>

Author: 3 people