The PR that added the warning (https://github.com/pytorch/pytorch/pull/13339 91) has a little discussion about its origin. As best I can tell, that PR’s author thought the function might be implemented via a different package in the future. I don’t think you have to worry about gumbel_softmax going away.