LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
┏━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━┳━━━━━━━┓
┃   ┃ Name                   ┃ Type                   ┃ Params ┃ Mode  ┃ FLOPs ┃
┡━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━╇━━━━━━━┩
│ 0 │ embedding              │ MultiInputVectorEmbed… │  220 K │ train │     0 │
│ 1 │ encoder                │ JetEncoder             │  332 K │ train │     0 │
│ 2 │ branch_decoders        │ ModuleList             │  9.6 M │ train │     0 │
│ 3 │ regression_decoder     │ RegressionDecoder      │      0 │ train │     0 │
│ 4 │ classification_decoder │ ClassificationDecoder  │      0 │ train │     0 │
│   │ other params           │ n/a                    │     48 │ n/a   │   n/a │
└───┴────────────────────────┴────────────────────────┴────────┴───────┴───────┘
Trainable params: 10.1 M
Non-trainable params: 86
Total params: 10.1 M
Total estimated model params size (MB): 40.531
Modules in train mode: 4351
Modules in eval mode: 0
Total FLOPs: 0
/opt/conda/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/opt/conda/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=5` in the `DataLoader` to improve performance.
/opt/conda/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/opt/conda/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=5` in the `DataLoader` to improve performance.
`Trainer.fit` stopped: `max_epochs=15` reached.
Epoch 14/14 ━━━━━━━━━━━━━━━━━━━ 3023/3023 0:39:21 • 0:00:00 1.29it/s v_num: 2uix