Support return embedding for MLP layer

[JiaxuanYou]

The added return_emb layer is useful when user would like to extract embeddings along with predictions

[codecov]

Codecov Report

Merging #4625 (cb9c230) into master (e3ba9d3) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #4625   +/-   ##
=======================================
  Coverage   82.88%   82.88%           
=======================================
  Files         316      316           
  Lines       16675    16677    +2     
=======================================
+ Hits        13821    13823    +2     
  Misses       2854     2854           

Impacted Files	Coverage Δ
torch_geometric/nn/models/mlp.py	98.46% <100.00%> (+0.04%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update e3ba9d3...cb9c230. Read the comment docs.

[rusty1s]

Can we do this as part of forward?

def forward(self, x, return_emb=True)

[JiaxuanYou]

Can we do this as part of forward?

def forward(self, x, return_emb=True)

Sounds good!

[JiaxuanYou]

Updated