| ============================================================================================================== |
| Layer (type:depth-idx) Output Shape Param # |
| ============================================================================================================== |
| TrainingWrapper -- -- |
| ├─FimOdeon: 1-1 -- -- |
| │ └─TrajectoryEncoder: 2-1 -- 896 |
| │ │ └─TransformerEncoder: 3-1 [1, 1194, 256] 1,579,520 |
| │ └─Sequential: 2-2 -- -- |
| │ │ └─Linear: 3-2 [1, 2400, 256] 1,024 |
| │ │ └─ReLU: 3-3 [1, 2400, 256] -- |
| │ │ └─Linear: 3-4 [1, 2400, 256] 65,792 |
| │ └─AttentionOperator: 2-3 -- -- |
| │ │ └─ModuleList: 3-5 -- 6,318,080 |
| │ │ └─MLP: 3-6 [1, 2400, 3] 132,355 |
| ├─UncertaintyEstimator: 1-2 -- -- |
| │ └─AttentionOperator: 2-4 -- -- |
| │ │ └─ModuleList: 3-7 -- 4,738,560 |
| │ │ └─MLP: 3-8 [1, 2400, 1] 131,841 |
| ============================================================================================================== |
| Total params: 12,968,068 |
| Trainable params: 12,968,068 |
| Non-trainable params: 0 |
| Total mult-adds (Units.MEGABYTES): 12.97 |
| ============================================================================================================== |
| Input size (MB): 0.09 |
| Forward/backward pass size (MB): 771.15 |
| Params size (MB): 51.87 |
| Estimated Total Size (MB): 823.11 |
| ============================================================================================================== |
