Pawel Gniewek

No, it's not about PyTorch Lightning integration? Simply, when I save the checkpoint, PNA conv layers don't save this dictionary ``` self.avg_deg: Dict[str, float] = { 'lin': float((bin_degrees * deg).sum())...

So rather than passing a `deg` tensor, mybe it's better to pass two values `avg_deg_lin` and `avg_deg_log` and then treat them as either trainable, or registered parameters.

I have this kind of an idea in mind, but requires changes to `pna_conv` so I'm not sure if it's a desirable approach: ``` def __init__( self, aggr: Union[str, List[str],...

Right, this approach requires calculating the stats prior the object creation. I guess I can live with that :) but having trainable `deltas` would be nice ...

Yeah, I will send a PR. thx.