Fix MLP.jittable() bug in case return_emb=True (#4645)

* fix MLP

* changelog

CHANGELOG.md

@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for graph-level outputs in `to_hetero` ([#4582](https://github.com/pyg-team/pytorch_geometric/pull/4582))
 - Added `CHANGELOG.md` ([#4581](https://github.com/pyg-team/pytorch_geometric/pull/4581))
 ### Changed
+- Fixed `MLP.jittable()` bug in case `return_emb=True` ([#4645](https://github.com/pyg-team/pytorch_geometric/pull/4645))
 - The generated node features of `StochasticBlockModelDataset` are now ordered with respect to their labels ([#4617](https://github.com/pyg-team/pytorch_geometric/pull/4617))
 - Removed unnecessary colons and fixed typos in the documentation ([#4616](https://github.com/pyg-team/pytorch_geometric/pull/4616))
 - The `bias` argument in `TAGConv` is now actually applied ([#4597](https://github.com/pyg-team/pytorch_geometric/pull/4597))

torch_geometric/nn/models/mlp.py

@@ -137,7 +137,7 @@ def reset_parameters(self):
             if hasattr(norm, 'reset_parameters'):
                 norm.reset_parameters()
 
-    def forward(self, x: Tensor, return_emb: bool = False) -> Tensor:
+    def forward(self, x: Tensor, return_emb=None) -> Tensor:
         """"""
         x = self.lins[0](x)
         emb = x
@@ -151,7 +151,7 @@ def forward(self, x: Tensor, return_emb: bool = False) -> Tensor:
             x = F.dropout(x, p=self.dropout, training=self.training)
             x = lin.forward(x)
 
-        return (x, emb) if return_emb else x
+        return (x, emb) if isinstance(return_emb, bool) else x
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}({str(self.channel_list)[1:-1]})'