Hi,
Why do you convert your mask to bool() since you multiply it with a float Tensor later?
Where do you get the error? What is the full stack trace?
Can you print the size of some Tensors to check they are what you expect.