you can use torch.gt() or torch.le() . it results in a tensor with the same dimension as the input tensor then you can do what you want.
you can use torch.gt() or torch.le() . it results in a tensor with the same dimension as the input tensor then you can do what you want.