Comment by spandan-madan

Yup! That version was in Keras. It's now been re-written in PyTorch as well! Thanks to https://github.com/AnshulBasia.