LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
┏━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━┳━━━━━━━┓
┃   ┃ Name                   ┃ Type                   ┃ Params ┃ Mode  ┃ FLOPs ┃
┡━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━╇━━━━━━━┩
│ 0 │ embedding              │ MultiInputVectorEmbed… │  220 K │ train │     0 │
│ 1 │ encoder                │ JetEncoder             │  498 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.3 M
Non-trainable params: 86
Total params: 10.3 M
Total estimated model params size (MB): 41.196
Modules in train mode: 4413
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 'num_encoder_layers' was locked by 'sweep' (ignored update).
wandb: WARNING Config item 'num_branch_encoder_layers' was locked by 'sweep' (ignored update).
wandb: WARNING Config item 'num_attention_heads' was locked by 'sweep' (ignored update).
wandb: WARNING Config item 'dataset_limit' 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).
wandb: WARNING Config item 'epochs' 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=10` reached.
Epoch 9/9  ━━━━━━━━━━━━━━━━━━━━ 1209/1209 0:14:20 • 0:00:00 1.41it/s v_num: xc7d