Masked argmax in Pytorch?

[yg-smile (Smile)]

Thanks. I need to convert m to tensor then.