LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0] ┏━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━┳━━━━━━━┓ ┃ ┃ Name ┃ Type ┃ Params ┃ Mode ┃ FLOPs ┃ ┡━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━╇━━━━━━━┩ │ 0 │ embedding │ MultiInputVectorEmbed… │ 3.4 M │ train │ 0 │ │ 1 │ encoder │ JetEncoder │ 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.922 Modules in train mode: 2299 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. Epoch 0/49 0/400 0:00:00 • -:--:-- 0.00it/s v_num: d7rx Traceback (most recent call last): File "", line 198, in _run_module_as_main File "", line 88, in _run_code File "/opt/conda/lib/python3.12/site-packages/spanet/train.py", line 277, in main(**parser.parse_args().__dict__) File "/opt/conda/lib/python3.12/site-packages/spanet/train.py", line 207, in main trainer.fit(model, ckpt_path=checkpoint) File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 584, in fit call._call_and_handle_interrupt( File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/trainer/call.py", line 49, in _call_and_handle_interrupt return trainer_fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 630, in _fit_impl self._run(model, ckpt_path=ckpt_path, weights_only=weights_only) File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1079, in _run results = self._run_stage() ^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1123, in _run_stage self.fit_loop.run() File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/loops/fit_loop.py", line 217, in run self.advance() File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/loops/fit_loop.py", line 469, in advance self.epoch_loop.run(self._data_fetcher) File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 153, in run self.advance(data_fetcher) File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 352, in advance batch_output = self.automatic_optimization.run(trainer.optimizers[0], batch_idx, kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 192, in run self._optimizer_step(batch_idx, closure) File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 270, in _optimizer_step call._call_lightning_module_hook( File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/trainer/call.py", line 177, in _call_lightning_module_hook output = fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/core/module.py", line 1368, in optimizer_step optimizer.step(closure=optimizer_closure) File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/core/optimizer.py", line 154, in step step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/strategies/strategy.py", line 239, in optimizer_step return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/plugins/precision/precision.py", line 123, in optimizer_step return optimizer.step(closure=closure, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/optim/lr_scheduler.py", line 166, in wrapper return func.__get__(opt, opt.__class__)(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/optim/optimizer.py", line 530, in wrapper out = func(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/optim/optimizer.py", line 80, in _use_grad ret = func(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/optim/adam.py", line 227, in step loss = closure() ^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/plugins/precision/precision.py", line 109, in _wrap_closure closure_result = closure() ^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 146, in __call__ self._result = self.closure(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context return func(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 131, in closure step_output = self._step_fn() ^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 319, in _training_step training_step_output = call._call_strategy_hook(trainer, "training_step", *kwargs.values()) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/trainer/call.py", line 329, in _call_strategy_hook output = fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/pytorch_lightning/strategies/strategy.py", line 391, in training_step return self.lightning_module.training_step(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/spanet/network/jet_reconstruction/jet_reconstruction_training.py", line 214, in training_step outputs = self.forward(batch.sources) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/spanet/network/jet_reconstruction/jet_reconstruction_network.py", line 148, in forward ) = decoder(hidden, padding_masks, sequence_masks, global_masks) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl return self._call_impl(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl return forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/spanet/network/layers/branch_decoder.py", line 128, in forward encoded_vectors, particle_vector, _ = self.encoder(event_vectors, padding_mask, sequence_mask) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl return self._call_impl(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl return forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/spanet/network/layers/stacked_encoder.py", line 49, in forward encoded_vectors = self.embedding(encoded_vectors, sequence_mask) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl return self._call_impl(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl return forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/spanet/network/layers/linear_stack.py", line 34, in forward output = layer(output, sequence_mask) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl return self._call_impl(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl return forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/spanet/network/layers/linear_block/gru_block.py", line 104, in forward output = self.linear_1(output) ^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl return self._call_impl(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl return forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/container.py", line 253, in forward input = module(input) ^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl return self._call_impl(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl return forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/dropout.py", line 73, in forward return F.dropout(input, self.p, self.training, self.inplace) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.12/site-packages/torch/nn/functional.py", line 1491, in dropout _VF.dropout_(input, p, training) if inplace else _VF.dropout(input, p, training) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 52.00 MiB. GPU 0 has a total capacity of 39.49 GiB of which 48.56 MiB is free. Including non-PyTorch memory, this process has 39.44 GiB memory in use. Of the allocated memory 38.92 GiB is allocated by PyTorch, and 27.45 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)