Does PyTorch provide a factorized output layer in nn module? It computes the gradient much faster in sparse output space. This will benefit a lot in NLP experiments.

I can propose such a layer if there isn't one.