in model_summary: also print the input shape, and output shape of each layer, e.g. for sanity check?

[mw66]

Description

Currently:

model_summary.py:104 -            summarize()]                                                                                                                                                      
  | Name                 | Type          | Params | Mode                                                                                                                                                                     
---------------------------------------------------------------                                                                                                                                                              
0 | loss                 | MAE           | 0      | train                                                                                                                                                                    
1 | padder_train         | ConstantPad1d | 0      | train                                                                                                                                                                    
2 | scaler               | TemporalNorm  | 0      | train                                                                                                                                                                    
3 | hist_exog_projection | MLPResidual   | 3.1 K  | train                                                                                                                                                                    
4 | dense_encoder        | Sequential    | 44.7 K | train                                                                                                                                                                    
5 | dense_decoder        | Sequential    | 8.4 K  | train                                                                                                                                                                    
6 | temporal_decoder     | MLPResidual   | 4.4 K  | train                                                                                                                                                                    
7 | global_skip          | Linear        | 64     | train                                                                                                                                                                    
---------------------------------------------------------------                                                                                                                                                              
60.6 K    Trainable params                                                                                                                                                                                                   
0         Non-trainable params                                                                                                                                                                                               
60.6 K    Total params                                                                                                                                                                                                       
0.243     Total estimated model params size (MB)                                                                                                                                                                             
30        Modules in train mode                                                                                                                                                                                              
0         Modules in eval mode

Can we also print the input shape, and output shape of each layer?
Thanks.

Use case

For final sanity check when the (long) training starts.