LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
┏━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━┳━━━━━━━┓
┃   ┃ Name                   ┃ Type                   ┃ Params ┃ Mode  ┃ FLOPs ┃
┡━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━╇━━━━━━━┩
│ 0 │ embedding              │ MultiInputVectorEmbed… │  3.4 M │ train │     0 │
│ 1 │ encoder                │ JetEncoderWithPairwise │  5.3 M │ train │     0 │
│ 2 │ branch_decoders        │ ModuleList             │ 75.8 M │ train │     0 │
│ 3 │ regression_decoder     │ RegressionDecoder      │      0 │ train │     0 │
│ 4 │ classification_decoder │ ClassificationDecoder  │      0 │ train │     0 │
│   │ other params           │ n/a                    │     12 │ n/a   │   n/a │
└───┴────────────────────────┴────────────────────────┴────────┴───────┴───────┘
Trainable params: 84.5 M
Non-trainable params: 50
Total params: 84.5 M
Total estimated model params size (MB): 337.925
Modules in train mode: 2326
Modules in eval mode: 0
Total FLOPs: 0
wandb: WARNING Config item 'hidden_dim' was locked by 'sweep' (ignored update).
wandb: WARNING Config item 'num_embedding_layers' was locked by 'sweep' (ignored update).
wandb: WARNING Config item 'batch_size' was locked by 'sweep' (ignored update).
wandb: WARNING Config item 'learning_rate' was locked by 'sweep' (ignored update).
wandb: WARNING Config item 'dropout' was locked by 'sweep' (ignored update).
/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=3` 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=3` in the `DataLoader` to improve performance.
`Trainer.fit` stopped: `max_epochs=50` reached.
Epoch 49/49 ━━━━━━━━━━━━━━━━━━━━━ 801/801 0:14:26 • 0:00:00 0.92it/s v_num: teh2