You probably want to upgrade your PyTorch version, I did check that the snippet above works with recent PyTorch.
Boolean tensors have been introduced in PyTorch 1.2 or so and PyTorch 1.3 and following refined the semantics.
Best regards
Thomas
P.S.: We don’t use at::Tensor but torch::Tensor these days .