With the latest 0.4 release, you can use the .to
function along with device
objects to make this much cleaner.
# on top of your script somewhere
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# in your code, use .to(device)
(x_u == 0).to(device, dtype=torch.float32)