add v2.0.0

clic/clusters/v2.0.0/README.md

@@ -0,0 +1,145 @@
+# Model Card for mlpf-clic-clusters-v2.0.0
+
+This model reconstructs particles in a detector, based on the tracks and calorimeter clusters recorded by the detector.
+
+## Model Details
+
+### Model Description
+
+- **Developed by:** Joosep Pata, Eric Wulff, Farouk Mokhtar, Mengke Zhang, David Southwick, Maria Girone, David Southwick, Javier Duarte, Michael Kagan
+- **Model type:** transformer
+- **License:** Apache License
+
+### Model Sources
+
+- **Repository:** https://github.com/jpata/particleflow/releases/tag/v2.0.0
+
+## Uses
+### Direct Use
+
+This model may be used to study the physics and computational performance on ML-based reconstruction in simulation.
+
+### Out-of-Scope Use
+
+This model is not intended for physics measurements on real data. 
+
+## Bias, Risks, and Limitations
+
+The model has only been trained on simulation data and has not been validated against real data.
+The model has not been peer reviewed or published in a peer-reviewed journal.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+```
+#get the code
+git clone https://github.com/jpata/particleflow
+cd particleflow
+git checkout v2.0.0
+
+#get the models
+git clone https://huggingface.co/jpata/particleflow models
+```
+
+## Training Details
+Trained on 8x MI250X for 26 epochs over ~3 days.
+The training was continued twice from a checkpoint due to the 24h time limit.
+
+### Training Data
+The following datasets were used:
+```
+/eos/user/j/jpata/mlpf/tensorflow_datasets/clic/clic_edm_qq_pf/2.3.0
+/eos/user/j/jpata/mlpf/tensorflow_datasets/clic/clic_edm_ttbar_pf/2.3.0
+/eos/user/j/jpata/mlpf/tensorflow_datasets/clic/clic_edm_ww_fullhad_pf/2.3.0
+```
+
+The truth and target definition was updated in [jpata/particleflow#352](https://github.com/jpata/particleflow/pull/352) have an updated truth and target definition with respect to [Pata, J., Wulff, E., Mokhtar, F. et al. Improved particle-flow event reconstruction with scalable neural networks for current and future particle detectors. Commun Phys 7, 124 (2024)](https://doi.org/10.1038/s42005-024-01599-5).
+
+In particular, target particles for MLPF reconstruction are based on status=1 particles.
+For non-interacting status=1, the direct children interacting status=0 are used instead.
+
+The datasets were generated using Key4HEP with the following scripts:
+- https://github.com/HEP-KBFI/key4hep-sim/releases/tag/v1.0.0
+- https://github.com/HEP-KBFI/key4hep-sim/blob/v1.0.0/clic/run_sim.sh
+
+## Training Procedure 
+
+```bash
+#!/bin/bash
+#SBATCH --job-name=mlpf-train
+#SBATCH --account=project_465000301
+#SBATCH --time=3-00:00:00
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=32
+#SBATCH --mem=200G
+#SBATCH --gpus-per-task=8
+#SBATCH --partition=small-g
+#SBATCH --no-requeue
+#SBATCH -o logs/slurm-%x-%j-%N.out
+
+cd /scratch/project_465000301/particleflow
+
+module load LUMI/24.03 partition/G
+
+export IMG=/scratch/project_465000301/pytorch-rocm6.2.simg
+export PYTHONPATH=`pwd`
+export TFDS_DATA_DIR=/scratch/project_465000301/tensorflow_datasets
+#export MIOPEN_DISABLE_CACHE=true
+export MIOPEN_USER_DB_PATH=/tmp/${USER}-${SLURM_JOB_ID}-miopen-cache
+export MIOPEN_CUSTOM_CACHE_DIR=${MIOPEN_USER_DB_PATH}
+export TF_CPP_MAX_VLOG_LEVEL=-1 #to suppress ROCm fusion is enabled messages
+export ROCM_PATH=/opt/rocm
+#export NCCL_DEBUG=INFO
+#export MIOPEN_ENABLE_LOGGING=1
+#export MIOPEN_ENABLE_LOGGING_CMD=1
+#export MIOPEN_LOG_LEVEL=4
+export KERAS_BACKEND=torch
+
+env
+
+#TF training
+singularity exec \
+    --rocm \
+    -B /scratch/project_465000301 \
+    -B /tmp \
+    --env LD_LIBRARY_PATH=/opt/rocm/lib/ \
+    --env CUDA_VISIBLE_DEVICES=$ROCR_VISIBLE_DEVICES \
+     $IMG python3 mlpf/pipeline.py --dataset clic --gpus 8 \
+     --data-dir $TFDS_DATA_DIR --config parameters/pytorch/pyg-clic.yaml \
+     --train --gpu-batch-multiplier 128 --num-workers 8 --prefetch-factor 100 --checkpoint-freq 1 --conv-type attention --dtype bfloat16 --lr 0.0001 --num-epochs 30
+```
+
+## Evaluation
+```bash
+#!/bin/bash
+#SBATCH --partition gpu
+#SBATCH --gres gpu:mig:1
+#SBATCH --mem-per-gpu 200G
+#SBATCH -o logs/slurm-%x-%j-%N.out
+
+IMG=/home/software/singularity/pytorch.simg:2024-08-18
+cd ~/particleflow
+
+WEIGHTS=models/clic/clusters/v2.0.0/checkpoints/checkpoint-29-1.901667.pth
+singularity exec -B /scratch/persistent --nv \
+     --env PYTHONPATH=`pwd` \
+     --env KERAS_BACKEND=torch \
+     $IMG  python3 mlpf/pyg_pipeline.py --dataset clic --gpus 1 \
+     --data-dir /scratch/persistent/joosep/tensorflow_datasets --config parameters/pytorch/pyg-clic.yaml \
+     --test --make-plots --gpu-batch-multiplier 100 --load $WEIGHTS --dtype bfloat16 --prefetch-factor 10 --num-workers 8 --load $WEIGHTS
+~               
+```
+
+## Citation
+
+## Glossary
+
+- PF: particle flow reconstruction
+- MLPF: machine learning for particle flow
+- CLIC: Compact Linear Collider
+
+## Model Card Contact
+
+Joosep Pata, [email protected]

clic/clusters/v2.0.0/history/epoch_1.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.22410098989764854, "Regression_eta": 0.00027124753357291037, "Regression_sin_phi": 0.00018932005113030378, "Regression_cos_phi": 0.00018533344456782053, "Regression_energy": 0.2203152060866305, "Classification_binary": 2.9749025808813028, "Classification": 0.029067417771931017, "MET": 7.935175555277818, "Sliced_Wasserstein_Loss": 47.84886087701757, "Total": 3.449032113715898}, "valid": {"Regression_pt": 0.1809588171661475, "Regression_eta": 0.00020314403285550918, "Regression_sin_phi": 0.00014497815029414137, "Regression_cos_phi": 0.00014624492148494414, "Regression_energy": 0.17771335614072548, "Classification_binary": 2.486452471864952, "Classification": 0.022163982330027882, "MET": 4.977831089228296, "Sliced_Wasserstein_Loss": 42.26595156752411, "Total": 2.867782857717042}, "epoch_train_time": 11350.586619138718, "epoch_valid_time": 420.86845803260803, "epoch_total_time": 11777.652925252914}

clic/clusters/v2.0.0/history/epoch_10.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.11755432991994447, "Regression_eta": 0.0001295076287009004, "Regression_sin_phi": 8.343577538195924e-05, "Regression_cos_phi": 8.436581557554069e-05, "Regression_energy": 0.1157861150276523, "Classification_binary": 1.7303218680813455, "Classification": 0.011502766537676538, "MET": 4.10354491769033, "Sliced_Wasserstein_Loss": 34.71469923939437, "Total": 1.9754638846236252}, "valid": {"Regression_pt": 0.11910374877537179, "Regression_eta": 0.0001345980014065071, "Regression_sin_phi": 8.7266393796424e-05, "Regression_cos_phi": 8.84002236308009e-05, "Regression_energy": 0.1177366103390022, "Classification_binary": 1.7656793232516077, "Classification": 0.011443455196270222, "MET": 4.119582747186495, "Sliced_Wasserstein_Loss": 34.6036349477492, "Total": 2.0142735128617364}, "epoch_train_time": 11209.789572238922, "epoch_valid_time": 416.5914318561554, "epoch_total_time": 11632.885925292969}

clic/clusters/v2.0.0/history/epoch_11.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.11599312542404656, "Regression_eta": 0.0001275183183604386, "Regression_sin_phi": 8.20733652576653e-05, "Regression_cos_phi": 8.309921048876388e-05, "Regression_energy": 0.11427387117130589, "Classification_binary": 1.7118472472682473, "Classification": 0.01130196340185219, "MET": 4.07779747982431, "Sliced_Wasserstein_Loss": 34.498703310241396, "Total": 1.953710853806599}, "valid": {"Regression_pt": 0.11830525781564007, "Regression_eta": 0.00013354199300624935, "Regression_sin_phi": 8.635481645823291e-05, "Regression_cos_phi": 8.725384040660797e-05, "Regression_energy": 0.11687312969440816, "Classification_binary": 1.75138744096664, "Classification": 0.011257291836753919, "MET": 4.050699608118971, "Sliced_Wasserstein_Loss": 34.49408410369775, "Total": 1.998130714931672}, "epoch_train_time": 11221.056280136108, "epoch_valid_time": 417.03449153900146, "epoch_total_time": 11644.654552698135}

clic/clusters/v2.0.0/history/epoch_12.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.1146443888317294, "Regression_eta": 0.00012577207509729968, "Regression_sin_phi": 8.092914508285054e-05, "Regression_cos_phi": 8.200626852784459e-05, "Regression_energy": 0.11298052914873769, "Classification_binary": 1.6960452357252536, "Classification": 0.011135802577519617, "MET": 4.0526077189865735, "Sliced_Wasserstein_Loss": 34.313254356520495, "Total": 1.9350961860359235}, "valid": {"Regression_pt": 0.11705597021955386, "Regression_eta": 0.00013250527466225086, "Regression_sin_phi": 8.490811973522714e-05, "Regression_cos_phi": 8.627060333632197e-05, "Regression_energy": 0.11595015449155949, "Classification_binary": 1.7400233621382637, "Classification": 0.011120271836063103, "MET": 4.019951705687299, "Sliced_Wasserstein_Loss": 34.405772709003216, "Total": 1.9844535018086817}, "epoch_train_time": 11210.498413801193, "epoch_valid_time": 415.2183647155762, "epoch_total_time": 11632.256944417953}

clic/clusters/v2.0.0/history/epoch_13.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.11343557714411334, "Regression_eta": 0.00012412126226742562, "Regression_sin_phi": 7.985220042624962e-05, "Regression_cos_phi": 8.097113481403641e-05, "Regression_energy": 0.11182199452812544, "Classification_binary": 1.6819292277442508, "Classification": 0.010991795440006896, "MET": 4.028762911101985, "Sliced_Wasserstein_Loss": 34.144026299814314, "Total": 1.9184633779995715}, "valid": {"Regression_pt": 0.11587548945877713, "Regression_eta": 0.000130290023000294, "Regression_sin_phi": 8.379978957283535e-05, "Regression_cos_phi": 8.549349101026725e-05, "Regression_energy": 0.11482585158762058, "Classification_binary": 1.7307191707696945, "Classification": 0.011027215301415545, "MET": 4.010250452170418, "Sliced_Wasserstein_Loss": 34.209244372990355, "Total": 1.9727471550944533}, "epoch_train_time": 11217.373317480087, "epoch_valid_time": 417.2627513408661, "epoch_total_time": 11641.021118164062}

clic/clusters/v2.0.0/history/epoch_14.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.11233104220323437, "Regression_eta": 0.00012261849474079386, "Regression_sin_phi": 7.893113370725656e-05, "Regression_cos_phi": 8.002588057957313e-05, "Regression_energy": 0.11076830332831916, "Classification_binary": 1.6691825832916727, "Classification": 0.010866083156856499, "MET": 4.006217861734038, "Sliced_Wasserstein_Loss": 33.99698257391801, "Total": 1.903428974432224}, "valid": {"Regression_pt": 0.11487826135764168, "Regression_eta": 0.0001287383186088881, "Regression_sin_phi": 8.293907742025001e-05, "Regression_cos_phi": 8.436576538147267e-05, "Regression_energy": 0.1137227539847518, "Classification_binary": 1.7209353333500803, "Classification": 0.01087729647228572, "MET": 3.97639042403537, "Sliced_Wasserstein_Loss": 34.00300944533762, "Total": 1.9607107867765274}, "epoch_train_time": 11210.823578119278, "epoch_valid_time": 416.7164807319641, "epoch_total_time": 11634.004011631012}

clic/clusters/v2.0.0/history/epoch_15.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.11133088862584363, "Regression_eta": 0.00012126362778394329, "Regression_sin_phi": 7.814369854152654e-05, "Regression_cos_phi": 7.921281056376189e-05, "Regression_energy": 0.10981926375194169, "Classification_binary": 1.6575475266926154, "Classification": 0.010756359800511199, "MET": 3.9855149041208398, "Sliced_Wasserstein_Loss": 33.85943525924868, "Total": 1.8897301780549207}, "valid": {"Regression_pt": 0.11403750698666097, "Regression_eta": 0.00012722494517875254, "Regression_sin_phi": 8.216095890646195e-05, "Regression_cos_phi": 8.354090226041542e-05, "Regression_energy": 0.1130240388238545, "Classification_binary": 1.7124789615152733, "Classification": 0.010770042149583627, "MET": 4.015169375502412, "Sliced_Wasserstein_Loss": 33.934761856913184, "Total": 1.950602736887058}, "epoch_train_time": 11219.220667362213, "epoch_valid_time": 415.95088052749634, "epoch_total_time": 11641.601608276367}

clic/clusters/v2.0.0/history/epoch_16.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.11042405319050404, "Regression_eta": 0.0001200718328690907, "Regression_sin_phi": 7.746683974007916e-05, "Regression_cos_phi": 7.841118300511077e-05, "Regression_energy": 0.10896015970932724, "Classification_binary": 1.6468976809116556, "Classification": 0.010660380469851963, "MET": 3.9689274300099986, "Sliced_Wasserstein_Loss": 33.72894050849879, "Total": 1.8772180146139836}, "valid": {"Regression_pt": 0.11350140832244775, "Regression_eta": 0.00012682964755791177, "Regression_sin_phi": 8.166588579343446e-05, "Regression_cos_phi": 8.278518819348989e-05, "Regression_energy": 0.1124117639670418, "Classification_binary": 1.7058728773110932, "Classification": 0.010691288889796022, "MET": 3.9297477893890673, "Sliced_Wasserstein_Loss": 33.838901728295816, "Total": 1.9427690413987138}, "epoch_train_time": 11219.385544538498, "epoch_valid_time": 416.36424255371094, "epoch_total_time": 11641.983350753784}

clic/clusters/v2.0.0/history/epoch_17.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.10959823471088952, "Regression_eta": 0.0001189503613888681, "Regression_sin_phi": 7.682247785751589e-05, "Regression_cos_phi": 7.775076422209809e-05, "Regression_energy": 0.10818125312455364, "Classification_binary": 1.6371782267711756, "Classification": 0.010576099689441686, "MET": 3.9535671801349808, "Sliced_Wasserstein_Loss": 33.61642086844736, "Total": 1.8658073958184545}, "valid": {"Regression_pt": 0.1128469093052904, "Regression_eta": 0.00012563044426908832, "Regression_sin_phi": 8.105971997190518e-05, "Regression_cos_phi": 8.237407905112508e-05, "Regression_energy": 0.11182280414740756, "Classification_binary": 1.7019418207395498, "Classification": 0.01064330833901163, "MET": 3.93617269141881, "Sliced_Wasserstein_Loss": 33.74690514469454, "Total": 1.9375439610128617}, "epoch_train_time": 11235.242618560791, "epoch_valid_time": 417.24506163597107, "epoch_total_time": 11658.845057725906}

clic/clusters/v2.0.0/history/epoch_18.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.10883164674856717, "Regression_eta": 0.00011799402227403095, "Regression_sin_phi": 7.631180712434671e-05, "Regression_cos_phi": 7.714453534966894e-05, "Regression_energy": 0.1074518651353378, "Classification_binary": 1.6280019428028139, "Classification": 0.010498198351337643, "MET": 3.9387966897586058, "Sliced_Wasserstein_Loss": 33.506046011284106, "Total": 1.855053351753321}, "valid": {"Regression_pt": 0.11230542345445639, "Regression_eta": 0.00012436031911917437, "Regression_sin_phi": 8.059409461987363e-05, "Regression_cos_phi": 8.169540063361263e-05, "Regression_energy": 0.1112928973133541, "Classification_binary": 1.6954759407656752, "Classification": 0.010563235267565564, "MET": 3.913439823653537, "Sliced_Wasserstein_Loss": 33.69031601688103, "Total": 1.9299245754622187}, "epoch_train_time": 11231.461406707764, "epoch_valid_time": 416.85057735443115, "epoch_total_time": 11654.869730949402}

clic/clusters/v2.0.0/history/epoch_19.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.10812470777054171, "Regression_eta": 0.00011713886036223777, "Regression_sin_phi": 7.583019376192718e-05, "Regression_cos_phi": 7.664578384270822e-05, "Regression_energy": 0.10678898724568366, "Classification_binary": 1.6196017756391945, "Classification": 0.010433087500822984, "MET": 3.9279115260319952, "Sliced_Wasserstein_Loss": 33.40636605484931, "Total": 1.8452181217415369}, "valid": {"Regression_pt": 0.11187282942498995, "Regression_eta": 0.00012365884336244638, "Regression_sin_phi": 8.001324639826343e-05, "Regression_cos_phi": 8.111327790754018e-05, "Regression_energy": 0.11095981413911776, "Classification_binary": 1.6907678732918006, "Classification": 0.010488686760905471, "MET": 3.9330002763263665, "Sliced_Wasserstein_Loss": 33.609555868167206, "Total": 1.924373869573955}, "epoch_train_time": 11233.42175936699, "epoch_valid_time": 416.27726459503174, "epoch_total_time": 11656.115841150284}

clic/clusters/v2.0.0/history/epoch_2.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.16195459451663155, "Regression_eta": 0.00017737329150928528, "Regression_sin_phi": 0.00012810314370809053, "Regression_cos_phi": 0.00012819214627293583, "Regression_energy": 0.1588706147447686, "Classification_binary": 2.2734354076649765, "Classification": 0.01913727421264038, "MET": 4.808671305884873, "Sliced_Wasserstein_Loss": 40.444829310098555, "Total": 2.613829162798172}, "valid": {"Regression_pt": 0.15254765844805065, "Regression_eta": 0.00017308740370526574, "Regression_sin_phi": 0.00012081858621149584, "Regression_cos_phi": 0.0001220595798308443, "Regression_energy": 0.15012515150849076, "Classification_binary": 2.179849213725884, "Classification": 0.017035812426993317, "MET": 4.531436520297428, "Sliced_Wasserstein_Loss": 39.15094704581993, "Total": 2.4999742514067522}, "epoch_train_time": 11215.72525548935, "epoch_valid_time": 416.6873936653137, "epoch_total_time": 11638.765037059784}

clic/clusters/v2.0.0/history/epoch_20.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.10748206450417798, "Regression_eta": 0.00011632373853677205, "Regression_sin_phi": 7.539108095467388e-05, "Regression_cos_phi": 7.61319545145938e-05, "Regression_energy": 0.10618606347251731, "Classification_binary": 1.6118662735680618, "Classification": 0.010375622391206948, "MET": 3.9123275357984575, "Sliced_Wasserstein_Loss": 33.31785860948436, "Total": 1.836176584148693}, "valid": {"Regression_pt": 0.11135896639808582, "Regression_eta": 0.00012300748533757935, "Regression_sin_phi": 7.96666386809763e-05, "Regression_cos_phi": 8.063100733557698e-05, "Regression_energy": 0.11043660939698051, "Classification_binary": 1.6876306270096464, "Classification": 0.010441160048702523, "MET": 3.932484676446945, "Sliced_Wasserstein_Loss": 33.580541599678455, "Total": 1.9201511002813505}, "epoch_train_time": 11212.203944444656, "epoch_valid_time": 416.51384139060974, "epoch_total_time": 11635.304948091507}

clic/clusters/v2.0.0/history/epoch_21.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.10687463349264838, "Regression_eta": 0.00011559049328162557, "Regression_sin_phi": 7.500507681391236e-05, "Regression_cos_phi": 7.571664610482542e-05, "Regression_energy": 0.10561656475135249, "Classification_binary": 1.6046765919600772, "Classification": 0.0103237782524511, "MET": 3.9033355168011714, "Sliced_Wasserstein_Loss": 33.23323453792315, "Total": 1.8277580044368662}, "valid": {"Regression_pt": 0.11090413673131029, "Regression_eta": 0.00012243332586871084, "Regression_sin_phi": 7.941454074007138e-05, "Regression_cos_phi": 8.034869022308055e-05, "Regression_energy": 0.10994413811294212, "Classification_binary": 1.6844999748794212, "Classification": 0.010403036381270724, "MET": 3.921498191318328, "Sliced_Wasserstein_Loss": 33.49826919212219, "Total": 1.9160335234123793}, "epoch_train_time": 11223.394551753998, "epoch_valid_time": 416.8033883571625, "epoch_total_time": 11646.649324893951}

clic/clusters/v2.0.0/history/epoch_22.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.10632739198529674, "Regression_eta": 0.00011500313084153785, "Regression_sin_phi": 7.467548466804895e-05, "Regression_cos_phi": 7.536618679936418e-05, "Regression_energy": 0.10510465755171136, "Classification_binary": 1.597958801644408, "Classification": 0.01027920933693345, "MET": 3.8968947516426224, "Sliced_Wasserstein_Loss": 33.15556706184831, "Total": 1.8199361586737608}, "valid": {"Regression_pt": 0.11054084998618369, "Regression_eta": 0.00012211642464640823, "Regression_sin_phi": 7.891557224310479e-05, "Regression_cos_phi": 8.000904915800431e-05, "Regression_energy": 0.1094820764670418, "Classification_binary": 1.6824691330888264, "Classification": 0.010369256071722392, "MET": 3.902202446744373, "Sliced_Wasserstein_Loss": 33.42549236334405, "Total": 1.913142144795016}, "epoch_train_time": 11212.190004348755, "epoch_valid_time": 418.27385449409485, "epoch_total_time": 11636.898732662201}

clic/clusters/v2.0.0/history/epoch_23.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.10582694903797225, "Regression_eta": 0.00011444309748439956, "Regression_sin_phi": 7.434966410590723e-05, "Regression_cos_phi": 7.502420717197697e-05, "Regression_energy": 0.1046316939376964, "Classification_binary": 1.5917921323739466, "Classification": 0.010237994696000459, "MET": 3.88400568891944, "Sliced_Wasserstein_Loss": 33.08106654406513, "Total": 1.8127523356038424}, "valid": {"Regression_pt": 0.11020492112138264, "Regression_eta": 0.00012220931973104692, "Regression_sin_phi": 7.871877150520251e-05, "Regression_cos_phi": 7.96159651501769e-05, "Regression_energy": 0.10915503793207396, "Classification_binary": 1.6803984123794211, "Classification": 0.010336244834580989, "MET": 3.8741342820538587, "Sliced_Wasserstein_Loss": 33.40897307073955, "Total": 1.9103751130426045}, "epoch_train_time": 11492.488549947739, "epoch_valid_time": 414.8949770927429, "epoch_total_time": 11913.816793203354}

clic/clusters/v2.0.0/history/epoch_24.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.10537340926742608, "Regression_eta": 0.00011395173589768809, "Regression_sin_phi": 7.40716226133955e-05, "Regression_cos_phi": 7.472390226765438e-05, "Regression_energy": 0.10420991406919547, "Classification_binary": 1.5862389357323954, "Classification": 0.010202363974161615, "MET": 3.877836648335952, "Sliced_Wasserstein_Loss": 33.016834648621625, "Total": 1.80628994965005}, "valid": {"Regression_pt": 0.1099853515625, "Regression_eta": 0.00012166760742089373, "Regression_sin_phi": 7.83466976555214e-05, "Regression_cos_phi": 7.932854522846136e-05, "Regression_energy": 0.10893743091840837, "Classification_binary": 1.6789492689911576, "Classification": 0.010306545392493344, "MET": 3.879703828376206, "Sliced_Wasserstein_Loss": 33.311364549839226, "Total": 1.9084574708601285}, "epoch_train_time": 11487.443885326385, "epoch_valid_time": 415.8364679813385, "epoch_total_time": 11910.326171636581}

clic/clusters/v2.0.0/history/epoch_25.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.10497148914552386, "Regression_eta": 0.00011349009575290078, "Regression_sin_phi": 7.381505275552502e-05, "Regression_cos_phi": 7.444741027454975e-05, "Regression_energy": 0.10383741740856217, "Classification_binary": 1.5812325359770032, "Classification": 0.010170177976806884, "MET": 3.8715518318811597, "Sliced_Wasserstein_Loss": 32.96086719754321, "Total": 1.800472839817526}, "valid": {"Regression_pt": 0.10967475915645096, "Regression_eta": 0.00012127597998959458, "Regression_sin_phi": 7.82751289594595e-05, "Regression_cos_phi": 7.920638445489277e-05, "Regression_energy": 0.10866716881656953, "Classification_binary": 1.676336100783762, "Classification": 0.010273796645774718, "MET": 3.89210711414791, "Sliced_Wasserstein_Loss": 33.23907757234727, "Total": 1.9052304185088424}, "epoch_train_time": 11486.676162958145, "epoch_valid_time": 414.46273970603943, "epoch_total_time": 11907.428564786911}

clic/clusters/v2.0.0/history/epoch_26.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.1046210579013266, "Regression_eta": 0.00011314179042189007, "Regression_sin_phi": 7.366601985925403e-05, "Regression_cos_phi": 7.426204076580892e-05, "Regression_energy": 0.10350892080986199, "Classification_binary": 1.5768355636694757, "Classification": 0.01014353230823331, "MET": 3.866421705381374, "Sliced_Wasserstein_Loss": 32.901092254677906, "Total": 1.795371727030067}, "valid": {"Regression_pt": 0.1095309045920418, "Regression_eta": 0.00012065126198281045, "Regression_sin_phi": 7.814548885707303e-05, "Regression_cos_phi": 7.902628643336403e-05, "Regression_energy": 0.10850268790192925, "Classification_binary": 1.6756780986233923, "Classification": 0.010261604946909227, "MET": 3.8605176723271706, "Sliced_Wasserstein_Loss": 33.22721563504823, "Total": 1.9042516579581994}, "epoch_train_time": 11479.693927288055, "epoch_valid_time": 417.1463315486908, "epoch_total_time": 11903.40635561943}

clic/clusters/v2.0.0/history/epoch_27.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.10432675354413655, "Regression_eta": 0.00011282419412583219, "Regression_sin_phi": 7.350894999902532e-05, "Regression_cos_phi": 7.411791177974941e-05, "Regression_energy": 0.10323505159140212, "Classification_binary": 1.5729909678081702, "Classification": 0.010120879512329713, "MET": 3.8640849856270534, "Sliced_Wasserstein_Loss": 32.860613662333954, "Total": 1.7909338844450793}, "valid": {"Regression_pt": 0.1093830268099377, "Regression_eta": 0.0001203362389776101, "Regression_sin_phi": 7.803604437990587e-05, "Regression_cos_phi": 7.884450664090957e-05, "Regression_energy": 0.10836989227981812, "Classification_binary": 1.673815564710611, "Classification": 0.010237073744991585, "MET": 3.8567765901326365, "Sliced_Wasserstein_Loss": 33.18617363344052, "Total": 1.90208233897709}, "epoch_train_time": 11509.053350448608, "epoch_valid_time": 415.4282293319702, "epoch_total_time": 11930.967039823532}

clic/clusters/v2.0.0/history/epoch_28.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.10406909120404674, "Regression_eta": 0.00011259657947391531, "Regression_sin_phi": 7.335984899289029e-05, "Regression_cos_phi": 7.39567553276778e-05, "Regression_energy": 0.10299630744715041, "Classification_binary": 1.569789416244108, "Classification": 0.010104072177398206, "MET": 3.8581435910941293, "Sliced_Wasserstein_Loss": 32.8221950435652, "Total": 1.7872187343772319}, "valid": {"Regression_pt": 0.1093371915663937, "Regression_eta": 0.00012041222052558825, "Regression_sin_phi": 7.789004653979727e-05, "Regression_cos_phi": 7.877612899737343e-05, "Regression_energy": 0.10830188027532155, "Classification_binary": 1.6738309510651126, "Classification": 0.010229814167574671, "MET": 3.84474383289791, "Sliced_Wasserstein_Loss": 33.204220257234724, "Total": 1.9019774605606914}, "epoch_train_time": 11456.624339103699, "epoch_valid_time": 415.2877013683319, "epoch_total_time": 11878.801152706146}

clic/clusters/v2.0.0/history/epoch_29.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.10386441724842879, "Regression_eta": 0.00011236702369087305, "Regression_sin_phi": 7.330479068154013e-05, "Regression_cos_phi": 7.385174694478111e-05, "Regression_energy": 0.10280298440903532, "Classification_binary": 1.5671975711059134, "Classification": 0.010089120700042293, "MET": 3.8600258556813314, "Sliced_Wasserstein_Loss": 32.78943633052421, "Total": 1.7842155361734038}, "valid": {"Regression_pt": 0.10923993763816318, "Regression_eta": 0.00011995665705089016, "Regression_sin_phi": 7.786119290870102e-05, "Regression_cos_phi": 7.867006722753837e-05, "Regression_energy": 0.10822336865391378, "Classification_binary": 1.6737055051748393, "Classification": 0.010221332905760149, "MET": 3.8721824130827973, "Sliced_Wasserstein_Loss": 33.17918006430868, "Total": 1.9016669074055466}, "epoch_train_time": 11500.102574825287, "epoch_valid_time": 413.5199978351593, "epoch_total_time": 11920.152396202087}

clic/clusters/v2.0.0/history/epoch_3.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.14496229716904283, "Regression_eta": 0.00015978972548468591, "Regression_sin_phi": 0.00011169671195010733, "Regression_cos_phi": 0.00011186725904968462, "Regression_energy": 0.14217958748638587, "Classification_binary": 2.0907613086523353, "Classification": 0.016321110054520672, "MET": 4.529948846593344, "Sliced_Wasserstein_Loss": 38.47457506070561, "Total": 2.3946121536209115}, "valid": {"Regression_pt": 0.14203008424814106, "Regression_eta": 0.00016207206287568022, "Regression_sin_phi": 0.00011132922969830381, "Regression_cos_phi": 0.00011116718555953341, "Regression_energy": 0.13993038655860632, "Classification_binary": 2.070844899266479, "Classification": 0.015637052480814158, "MET": 4.390331717242765, "Sliced_Wasserstein_Loss": 37.99571442926045, "Total": 2.368827873794212}, "epoch_train_time": 11230.485923290253, "epoch_valid_time": 417.6828181743622, "epoch_total_time": 11654.763090133667}

clic/clusters/v2.0.0/history/epoch_4.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.13628292533265338, "Regression_eta": 0.00015098359683227104, "Regression_sin_phi": 0.00010357710518610169, "Regression_cos_phi": 0.00010353283377447021, "Regression_energy": 0.1336098365691062, "Classification_binary": 1.999462269898943, "Classification": 0.015236734543642202, "MET": 4.386986289905013, "Sliced_Wasserstein_Loss": 37.41231163405228, "Total": 2.2849514801456934}, "valid": {"Regression_pt": 0.13568469473786676, "Regression_eta": 0.00015768849581384198, "Regression_sin_phi": 0.00010537840545752424, "Regression_cos_phi": 0.00010523528723088108, "Regression_energy": 0.13362879124485028, "Classification_binary": 1.9973497789389068, "Classification": 0.014792787607076467, "MET": 4.334894681973473, "Sliced_Wasserstein_Loss": 36.97285470257235, "Total": 2.28182400522508}, "epoch_train_time": 11210.8334004879, "epoch_valid_time": 414.8488028049469, "epoch_total_time": 11632.07339334488}

clic/clusters/v2.0.0/history/epoch_5.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.13111573834876267, "Regression_eta": 0.00014555711368887173, "Regression_sin_phi": 9.84128535330219e-05, "Regression_cos_phi": 9.837286452909245e-05, "Regression_energy": 0.12852008341721272, "Classification_binary": 1.9302377841111984, "Classification": 0.014384665977544503, "MET": 4.301009454006571, "Sliced_Wasserstein_Loss": 36.69936660834166, "Total": 2.2046026069936437}, "valid": {"Regression_pt": 0.13089269901778536, "Regression_eta": 0.0001510359658305668, "Regression_sin_phi": 0.00010049001005301522, "Regression_cos_phi": 0.00010077041061744812, "Regression_energy": 0.12910330916524318, "Classification_binary": 1.9329823779139872, "Classification": 0.013855981366810691, "MET": 4.278978471663987, "Sliced_Wasserstein_Loss": 36.31261555466238, "Total": 2.2071882536173635}, "epoch_train_time": 11231.486518859863, "epoch_valid_time": 416.5859453678131, "epoch_total_time": 11654.359743595123}

clic/clusters/v2.0.0/history/epoch_6.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.12751388822244145, "Regression_eta": 0.00014165317547387726, "Regression_sin_phi": 9.496446814371541e-05, "Regression_cos_phi": 9.485025641543919e-05, "Regression_energy": 0.12504516828232126, "Classification_binary": 1.8715302110859162, "Classification": 0.013408395532641542, "MET": 4.240105487162548, "Sliced_Wasserstein_Loss": 36.15900407084702, "Total": 2.137831593254535}, "valid": {"Regression_pt": 0.12807342431169613, "Regression_eta": 0.00014617216932045303, "Regression_sin_phi": 9.77680614140256e-05, "Regression_cos_phi": 9.759754613281446e-05, "Regression_energy": 0.12617450481059084, "Classification_binary": 1.879413528687701, "Classification": 0.01291795712191959, "MET": 4.29762233721865, "Sliced_Wasserstein_Loss": 35.808601286173634, "Total": 2.1469214730707393}, "epoch_train_time": 11222.83605313301, "epoch_valid_time": 414.3109006881714, "epoch_total_time": 11644.033034801483}

clic/clusters/v2.0.0/history/epoch_7.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.124092284040338, "Regression_eta": 0.0001369230135255772, "Regression_sin_phi": 9.062163847852581e-05, "Regression_cos_phi": 9.13321180116141e-05, "Regression_energy": 0.12206584885194972, "Classification_binary": 1.8204286943383088, "Classification": 0.012620926720638682, "MET": 4.214292489465791, "Sliced_Wasserstein_Loss": 35.66058420225682, "Total": 2.079524074685759}, "valid": {"Regression_pt": 0.1250619673652281, "Regression_eta": 0.0001421002329737427, "Regression_sin_phi": 9.3659998136318e-05, "Regression_cos_phi": 9.405039897685649e-05, "Regression_energy": 0.12348332543081793, "Classification_binary": 1.8350366132435691, "Classification": 0.012317114090996157, "MET": 4.232725834003215, "Sliced_Wasserstein_Loss": 35.38281752411576, "Total": 2.096229087118167}, "epoch_train_time": 11242.381155490875, "epoch_valid_time": 417.21450686454773, "epoch_total_time": 11665.937296628952}

clic/clusters/v2.0.0/history/epoch_8.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.12145471266381588, "Regression_eta": 0.00013390145562679627, "Regression_sin_phi": 8.719075732427978e-05, "Regression_cos_phi": 8.812197241301605e-05, "Regression_energy": 0.11957195114704595, "Classification_binary": 1.7823741418636623, "Classification": 0.012110475437451458, "MET": 4.17112482814955, "Sliced_Wasserstein_Loss": 35.2700730252821, "Total": 2.0358210545814885}, "valid": {"Regression_pt": 0.12264579944671926, "Regression_eta": 0.00014085805109459486, "Regression_sin_phi": 9.00495186496011e-05, "Regression_cos_phi": 9.156061138754106e-05, "Regression_energy": 0.12115546223434988, "Classification_binary": 1.8063186105807878, "Classification": 0.011966087120522257, "MET": 4.201982013665595, "Sliced_Wasserstein_Loss": 35.172462821543405, "Total": 2.0624090949055467}, "epoch_train_time": 11220.988495588303, "epoch_valid_time": 416.6119010448456, "epoch_total_time": 11644.025322198868}

clic/clusters/v2.0.0/history/epoch_9.json

@@ -0,0 +1 @@
+{"train": {"Regression_pt": 0.11933376495880053, "Regression_eta": 0.00013156201801646047, "Regression_sin_phi": 8.509110049032651e-05, "Regression_cos_phi": 8.591650690526353e-05, "Regression_energy": 0.11751268373491287, "Classification_binary": 1.753052298332381, "Classification": 0.011758926101725981, "MET": 4.1335741099485785, "Sliced_Wasserstein_Loss": 34.96227101485502, "Total": 2.001957728137052}, "valid": {"Regression_pt": 0.12088643653599779, "Regression_eta": 0.0001371413948451591, "Regression_sin_phi": 8.864979651963212e-05, "Regression_cos_phi": 8.976404881554018e-05, "Regression_energy": 0.11950649249208702, "Classification_binary": 1.7831902507033761, "Classification": 0.01168908122268137, "MET": 4.149929662379421, "Sliced_Wasserstein_Loss": 34.88086314308681, "Total": 2.0355873819332797}, "epoch_train_time": 11223.379476547241, "epoch_valid_time": 416.62553310394287, "epoch_total_time": 11646.32419371605}

clic/clusters/v2.0.0/hyperparameters.json

@@ -0,0 +1 @@
+{"Num of mlpf parameters": 89388050, "config": "parameters/pytorch/pyg-clic.yaml", "prefix": null, "data_dir": "/scratch/project_465000301/tensorflow_datasets", "gpus": 8, "gpu_batch_multiplier": 128, "dataset": "clic", "num_workers": 8, "prefetch_factor": 100, "resume_training": null, "load": "experiments/pyg-clic_20241011_102451_167094/checkpoints/checkpoint-22-1.913142.pth", "train": true, "test": null, "num_epochs": 30, "patience": null, "lr": 0.0001, "conv_type": "attention", "num_convs": null, "make_plots": null, "export_onnx": null, "ntrain": null, "ntest": null, "nvalid": null, "val_freq": null, "checkpoint_freq": 1, "hpo": null, "ray_train": false, "local": null, "ray_cpus": null, "ray_gpus": null, "raytune_num_samples": null, "comet": false, "comet_offline": false, "comet_step_freq": null, "experiments_dir": null, "pipeline": null, "dtype": "bfloat16", "attention_type": null, "test_datasets": {"clic_edm_qq_pf": {"version": "2.3.0"}, "clic_edm_ttbar_pf": {"version": "2.3.0"}, "clic_edm_ww_fullhad_pf": {"version": "2.3.0"}}}

clic/clusters/v2.0.0/test-config.yaml

@@ -0,0 +1,126 @@
+backend: pytorch
+checkpoint_freq: null
+comet: false
+comet_name: particleflow-pt
+comet_offline: false
+comet_step_freq: 100
+config: parameters/pytorch/pyg-clic.yaml
+conv_type: attention
+data_dir: /scratch/persistent/joosep/tensorflow_datasets
+dataset: clic
+dtype: bfloat16
+gpu_batch_multiplier: 100
+gpus: 1
+load: experiments/pyg-clic_20241011_102451_167094/checkpoints/checkpoint-29-1.901667.pth
+lr: 0.0001
+lr_schedule: cosinedecay
+lr_schedule_config:
+  onecycle:
+    pct_start: 0.3
+make_plots: true
+model:
+  attention:
+    activation: gelu
+    attention_type: math
+    conv_type: attention
+    dropout_conv_id_ff: 0.0
+    dropout_conv_id_mha: 0.0
+    dropout_conv_reg_ff: 0.1
+    dropout_conv_reg_mha: 0.1
+    dropout_ff: 0.1
+    head_dim: 32
+    num_convs: 6
+    num_heads: 32
+    use_pre_layernorm: true
+  cos_phi_mode: linear
+  energy_mode: direct-elemtype-split
+  eta_mode: linear
+  gnn_lsh:
+    activation: elu
+    bin_size: 32
+    conv_type: gnn_lsh
+    distance_dim: 128
+    embedding_dim: 512
+    ffn_dist_hidden_dim: 128
+    ffn_dist_num_layers: 2
+    layernorm: true
+    max_num_bins: 200
+    num_convs: 8
+    num_node_messages: 2
+    width: 512
+  input_encoding: split
+  learned_representation_mode: last
+  mamba:
+    activation: elu
+    conv_type: mamba
+    d_conv: 4
+    d_state: 16
+    dropout: 0.0
+    embedding_dim: 128
+    expand: 2
+    num_convs: 2
+    num_heads: 2
+    width: 128
+  pt_mode: direct-elemtype-split
+  sin_phi_mode: linear
+  trainable: all
+ntest: null
+ntrain: null
+num_epochs: 30
+num_workers: 8
+nvalid: null
+patience: 20
+prefetch_factor: 10
+ray_train: false
+raytune:
+  asha:
+    brackets: 1
+    grace_period: 10
+    max_t: 200
+    reduction_factor: 4
+  default_metric: val_loss
+  default_mode: min
+  hyperband:
+    max_t: 200
+    reduction_factor: 4
+  hyperopt:
+    n_random_steps: 10
+  local_dir: null
+  nevergrad:
+    n_random_steps: 10
+  sched: null
+  search_alg: null
+save_attention: true
+sort_data: false
+test: true
+test_dataset:
+  clic_edm_qq_pf:
+    version: 2.3.0
+  clic_edm_ttbar_pf:
+    version: 2.3.0
+  clic_edm_ww_fullhad_pf:
+    version: 2.3.0
+test_datasets: []
+train_dataset:
+  clic:
+    physical:
+      batch_size: 1
+      samples:
+        clic_edm_qq_pf:
+          version: 2.3.0
+        clic_edm_ttbar_pf:
+          version: 2.3.0
+        clic_edm_ww_fullhad_pf:
+          version: 2.3.0
+val_freq: null
+valid_dataset:
+  clic:
+    physical:
+      batch_size: 1
+      samples:
+        clic_edm_qq_pf:
+          version: 2.3.0
+        clic_edm_ttbar_pf:
+          version: 2.3.0
+        clic_edm_ww_fullhad_pf:
+          version: 2.3.0

clic/clusters/v2.0.0/train-config.yaml

@@ -0,0 +1,125 @@
+backend: pytorch
+checkpoint_freq: 1
+comet: false
+comet_name: particleflow-pt
+comet_offline: false
+comet_step_freq: 100
+config: parameters/pytorch/pyg-clic.yaml
+conv_type: attention
+data_dir: /scratch/project_465000301/tensorflow_datasets
+dataset: clic
+dtype: bfloat16
+gpu_batch_multiplier: 128
+gpus: 8
+load: experiments/pyg-clic_20241011_102451_167094/checkpoints/checkpoint-22-1.913142.pth
+lr: 0.0001
+lr_schedule: cosinedecay
+lr_schedule_config:
+  onecycle:
+    pct_start: 0.3
+model:
+  attention:
+    activation: gelu
+    attention_type: math
+    conv_type: attention
+    dropout_conv_id_ff: 0.0
+    dropout_conv_id_mha: 0.0
+    dropout_conv_reg_ff: 0.1
+    dropout_conv_reg_mha: 0.1
+    dropout_ff: 0.1
+    head_dim: 32
+    num_convs: 6
+    num_heads: 32
+    use_pre_layernorm: true
+  cos_phi_mode: linear
+  energy_mode: direct-elemtype-split
+  eta_mode: linear
+  gnn_lsh:
+    activation: elu
+    bin_size: 32
+    conv_type: gnn_lsh
+    distance_dim: 128
+    embedding_dim: 512
+    ffn_dist_hidden_dim: 128
+    ffn_dist_num_layers: 2
+    layernorm: true
+    max_num_bins: 200
+    num_convs: 8
+    num_node_messages: 2
+    width: 512
+  input_encoding: split
+  learned_representation_mode: last
+  mamba:
+    activation: elu
+    conv_type: mamba
+    d_conv: 4
+    d_state: 16
+    dropout: 0.0
+    embedding_dim: 128
+    expand: 2
+    num_convs: 2
+    num_heads: 2
+    width: 128
+  pt_mode: direct-elemtype-split
+  sin_phi_mode: linear
+  trainable: all
+ntest: null
+ntrain: null
+num_epochs: 30
+num_workers: 8
+nvalid: null
+patience: 20
+prefetch_factor: 100
+ray_train: false
+raytune:
+  asha:
+    brackets: 1
+    grace_period: 10
+    max_t: 200
+    reduction_factor: 4
+  default_metric: val_loss
+  default_mode: min
+  hyperband:
+    max_t: 200
+    reduction_factor: 4
+  hyperopt:
+    n_random_steps: 10
+  local_dir: null
+  nevergrad:
+    n_random_steps: 10
+  sched: null
+  search_alg: null
+save_attention: true
+sort_data: false
+test_dataset:
+  clic_edm_qq_pf:
+    version: 2.3.0
+  clic_edm_ttbar_pf:
+    version: 2.3.0
+  clic_edm_ww_fullhad_pf:
+    version: 2.3.0
+test_datasets: []
+train: true
+train_dataset:
+  clic:
+    physical:
+      batch_size: 1
+      samples:
+        clic_edm_qq_pf:
+          version: 2.3.0
+        clic_edm_ttbar_pf:
+          version: 2.3.0
+        clic_edm_ww_fullhad_pf:
+          version: 2.3.0
+val_freq: null
+valid_dataset:
+  clic:
+    physical:
+      batch_size: 1
+      samples:
+        clic_edm_qq_pf:
+          version: 2.3.0
+        clic_edm_ttbar_pf:
+          version: 2.3.0
+        clic_edm_ww_fullhad_pf:
+          version: 2.3.0

clic/clusters/v2.0.0/train.log

@@ -0,0 +1,850 @@
+[2024-10-11 10:25:04,881] INFO: Will use torch.nn.parallel.DistributedDataParallel() and 8 gpus
+[2024-10-11 10:25:04,885] INFO: AMD Radeon Graphics
+[2024-10-11 10:25:04,885] INFO: AMD Radeon Graphics
+[2024-10-11 10:25:04,885] INFO: AMD Radeon Graphics
+[2024-10-11 10:25:04,885] INFO: AMD Radeon Graphics
+[2024-10-11 10:25:04,886] INFO: AMD Radeon Graphics
+[2024-10-11 10:25:04,886] INFO: AMD Radeon Graphics
+[2024-10-11 10:25:04,886] INFO: AMD Radeon Graphics
+[2024-10-11 10:25:04,886] INFO: AMD Radeon Graphics
+[2024-10-11 10:25:08,545] INFO: configured dtype=torch.bfloat16 for autocast
+[2024-10-11 10:25:10,626] INFO: using attention_type=math
+[2024-10-11 10:25:10,662] INFO: using attention_type=math
+[2024-10-11 10:25:10,697] INFO: using attention_type=math
+[2024-10-11 10:25:10,735] INFO: using attention_type=math
+[2024-10-11 10:25:10,770] INFO: using attention_type=math
+[2024-10-11 10:25:10,807] INFO: using attention_type=math
+[2024-10-11 10:25:10,843] INFO: using attention_type=math
+[2024-10-11 10:25:10,879] INFO: using attention_type=math
+[2024-10-11 10:25:10,914] INFO: using attention_type=math
+[2024-10-11 10:25:10,950] INFO: using attention_type=math
+[2024-10-11 10:25:10,987] INFO: using attention_type=math
+[2024-10-11 10:25:11,024] INFO: using attention_type=math
+[2024-10-11 10:25:16,061] INFO: DistributedDataParallel(
+  (module): MLPF(
+    (nn0_id): ModuleList(
+      (0-1): 2 x Sequential(
+        (0): Linear(in_features=17, out_features=1024, bias=True)
+        (1): GELU(approximate='none')
+        (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (3): Dropout(p=0.1, inplace=False)
+        (4): Linear(in_features=1024, out_features=1024, bias=True)
+      )
+    )
+    (nn0_reg): ModuleList(
+      (0-1): 2 x Sequential(
+        (0): Linear(in_features=17, out_features=1024, bias=True)
+        (1): GELU(approximate='none')
+        (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (3): Dropout(p=0.1, inplace=False)
+        (4): Linear(in_features=1024, out_features=1024, bias=True)
+      )
+    )
+    (conv_id): ModuleList(
+      (0-5): 6 x PreLnSelfAttentionLayer(
+        (mha): MultiheadAttention(
+          (out_proj): NonDynamicallyQuantizableLinear(in_features=1024, out_features=1024, bias=True)
+        )
+        (norm0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (seq): Sequential(
+          (0): Linear(in_features=1024, out_features=1024, bias=True)
+          (1): GELU(approximate='none')
+          (2): Linear(in_features=1024, out_features=1024, bias=True)
+          (3): GELU(approximate='none')
+        )
+        (dropout): Dropout(p=0.0, inplace=False)
+      )
+    )
+    (conv_reg): ModuleList(
+      (0-5): 6 x PreLnSelfAttentionLayer(
+        (mha): MultiheadAttention(
+          (out_proj): NonDynamicallyQuantizableLinear(in_features=1024, out_features=1024, bias=True)
+        )
+        (norm0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (seq): Sequential(
+          (0): Linear(in_features=1024, out_features=1024, bias=True)
+          (1): GELU(approximate='none')
+          (2): Linear(in_features=1024, out_features=1024, bias=True)
+          (3): GELU(approximate='none')
+        )
+        (dropout): Dropout(p=0.1, inplace=False)
+      )
+    )
+    (nn_binary_particle): Sequential(
+      (0): Linear(in_features=1024, out_features=1024, bias=True)
+      (1): GELU(approximate='none')
+      (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+      (3): Dropout(p=0.1, inplace=False)
+      (4): Linear(in_features=1024, out_features=2, bias=True)
+    )
+    (nn_pid): Sequential(
+      (0): Linear(in_features=1024, out_features=1024, bias=True)
+      (1): GELU(approximate='none')
+      (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+      (3): Dropout(p=0.1, inplace=False)
+      (4): Linear(in_features=1024, out_features=6, bias=True)
+    )
+    (nn_pt): RegressionOutput(
+      (nn): ModuleList(
+        (0-1): 2 x Sequential(
+          (0): Linear(in_features=1024, out_features=1024, bias=True)
+          (1): GELU(approximate='none')
+          (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+          (3): Dropout(p=0.1, inplace=False)
+          (4): Linear(in_features=1024, out_features=1, bias=True)
+        )
+      )
+    )
+    (nn_eta): RegressionOutput(
+      (nn): Sequential(
+        (0): Linear(in_features=1024, out_features=1024, bias=True)
+        (1): GELU(approximate='none')
+        (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (3): Dropout(p=0.1, inplace=False)
+        (4): Linear(in_features=1024, out_features=2, bias=True)
+      )
+    )
+    (nn_sin_phi): RegressionOutput(
+      (nn): Sequential(
+        (0): Linear(in_features=1024, out_features=1024, bias=True)
+        (1): GELU(approximate='none')
+        (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (3): Dropout(p=0.1, inplace=False)
+        (4): Linear(in_features=1024, out_features=2, bias=True)
+      )
+    )
+    (nn_cos_phi): RegressionOutput(
+      (nn): Sequential(
+        (0): Linear(in_features=1024, out_features=1024, bias=True)
+        (1): GELU(approximate='none')
+        (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (3): Dropout(p=0.1, inplace=False)
+        (4): Linear(in_features=1024, out_features=2, bias=True)
+      )
+    )
+    (nn_energy): RegressionOutput(
+      (nn): ModuleList(
+        (0-1): 2 x Sequential(
+          (0): Linear(in_features=1024, out_features=1024, bias=True)
+          (1): GELU(approximate='none')
+          (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+          (3): Dropout(p=0.1, inplace=False)
+          (4): Linear(in_features=1024, out_features=1, bias=True)
+        )
+      )
+    )
+    (final_norm_id): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+    (final_norm_reg): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+  )
+)
+[2024-10-11 10:25:16,062] INFO: Trainable parameters: 89388050
+[2024-10-11 10:25:16,062] INFO: Non-trainable parameters: 0
+[2024-10-11 10:25:16,062] INFO: Total parameters: 89388050
+[2024-10-11 10:25:16,067] INFO:                               Modules Trainable parameters Non-trainable parameters
+             module.nn0_id.0.0.weight                17408                        0
+               module.nn0_id.0.0.bias                 1024                        0
+             module.nn0_id.0.2.weight                 1024                        0
+               module.nn0_id.0.2.bias                 1024                        0
+             module.nn0_id.0.4.weight              1048576                        0
+               module.nn0_id.0.4.bias                 1024                        0
+             module.nn0_id.1.0.weight                17408                        0
+               module.nn0_id.1.0.bias                 1024                        0
+             module.nn0_id.1.2.weight                 1024                        0
+               module.nn0_id.1.2.bias                 1024                        0
+             module.nn0_id.1.4.weight              1048576                        0
+               module.nn0_id.1.4.bias                 1024                        0
+            module.nn0_reg.0.0.weight                17408                        0
+              module.nn0_reg.0.0.bias                 1024                        0
+            module.nn0_reg.0.2.weight                 1024                        0
+              module.nn0_reg.0.2.bias                 1024                        0
+            module.nn0_reg.0.4.weight              1048576                        0
+              module.nn0_reg.0.4.bias                 1024                        0
+            module.nn0_reg.1.0.weight                17408                        0
+              module.nn0_reg.1.0.bias                 1024                        0
+            module.nn0_reg.1.2.weight                 1024                        0
+              module.nn0_reg.1.2.bias                 1024                        0
+            module.nn0_reg.1.4.weight              1048576                        0
+              module.nn0_reg.1.4.bias                 1024                        0
+  module.conv_id.0.mha.in_proj_weight              3145728                        0
+    module.conv_id.0.mha.in_proj_bias                 3072                        0
+ module.conv_id.0.mha.out_proj.weight              1048576                        0
+   module.conv_id.0.mha.out_proj.bias                 1024                        0
+        module.conv_id.0.norm0.weight                 1024                        0
+          module.conv_id.0.norm0.bias                 1024                        0
+        module.conv_id.0.norm1.weight                 1024                        0
+          module.conv_id.0.norm1.bias                 1024                        0
+        module.conv_id.0.seq.0.weight              1048576                        0
+          module.conv_id.0.seq.0.bias                 1024                        0
+        module.conv_id.0.seq.2.weight              1048576                        0
+          module.conv_id.0.seq.2.bias                 1024                        0
+  module.conv_id.1.mha.in_proj_weight              3145728                        0
+    module.conv_id.1.mha.in_proj_bias                 3072                        0
+ module.conv_id.1.mha.out_proj.weight              1048576                        0
+   module.conv_id.1.mha.out_proj.bias                 1024                        0
+        module.conv_id.1.norm0.weight                 1024                        0
+          module.conv_id.1.norm0.bias                 1024                        0
+        module.conv_id.1.norm1.weight                 1024                        0
+          module.conv_id.1.norm1.bias                 1024                        0
+        module.conv_id.1.seq.0.weight              1048576                        0
+          module.conv_id.1.seq.0.bias                 1024                        0
+        module.conv_id.1.seq.2.weight              1048576                        0
+          module.conv_id.1.seq.2.bias                 1024                        0
+  module.conv_id.2.mha.in_proj_weight              3145728                        0
+    module.conv_id.2.mha.in_proj_bias                 3072                        0
+ module.conv_id.2.mha.out_proj.weight              1048576                        0
+   module.conv_id.2.mha.out_proj.bias                 1024                        0
+        module.conv_id.2.norm0.weight                 1024                        0
+          module.conv_id.2.norm0.bias                 1024                        0
+        module.conv_id.2.norm1.weight                 1024                        0
+          module.conv_id.2.norm1.bias                 1024                        0
+        module.conv_id.2.seq.0.weight              1048576                        0
+          module.conv_id.2.seq.0.bias                 1024                        0
+        module.conv_id.2.seq.2.weight              1048576                        0
+          module.conv_id.2.seq.2.bias                 1024                        0
+  module.conv_id.3.mha.in_proj_weight              3145728                        0
+    module.conv_id.3.mha.in_proj_bias                 3072                        0
+ module.conv_id.3.mha.out_proj.weight              1048576                        0
+   module.conv_id.3.mha.out_proj.bias                 1024                        0
+        module.conv_id.3.norm0.weight                 1024                        0
+          module.conv_id.3.norm0.bias                 1024                        0
+        module.conv_id.3.norm1.weight                 1024                        0
+          module.conv_id.3.norm1.bias                 1024                        0
+        module.conv_id.3.seq.0.weight              1048576                        0
+          module.conv_id.3.seq.0.bias                 1024                        0
+        module.conv_id.3.seq.2.weight              1048576                        0
+          module.conv_id.3.seq.2.bias                 1024                        0
+  module.conv_id.4.mha.in_proj_weight              3145728                        0
+    module.conv_id.4.mha.in_proj_bias                 3072                        0
+ module.conv_id.4.mha.out_proj.weight              1048576                        0
+   module.conv_id.4.mha.out_proj.bias                 1024                        0
+        module.conv_id.4.norm0.weight                 1024                        0
+          module.conv_id.4.norm0.bias                 1024                        0
+        module.conv_id.4.norm1.weight                 1024                        0
+          module.conv_id.4.norm1.bias                 1024                        0
+        module.conv_id.4.seq.0.weight              1048576                        0
+          module.conv_id.4.seq.0.bias                 1024                        0
+        module.conv_id.4.seq.2.weight              1048576                        0
+          module.conv_id.4.seq.2.bias                 1024                        0
+  module.conv_id.5.mha.in_proj_weight              3145728                        0
+    module.conv_id.5.mha.in_proj_bias                 3072                        0
+ module.conv_id.5.mha.out_proj.weight              1048576                        0
+   module.conv_id.5.mha.out_proj.bias                 1024                        0
+        module.conv_id.5.norm0.weight                 1024                        0
+          module.conv_id.5.norm0.bias                 1024                        0
+        module.conv_id.5.norm1.weight                 1024                        0
+          module.conv_id.5.norm1.bias                 1024                        0
+        module.conv_id.5.seq.0.weight              1048576                        0
+          module.conv_id.5.seq.0.bias                 1024                        0
+        module.conv_id.5.seq.2.weight              1048576                        0
+          module.conv_id.5.seq.2.bias                 1024                        0
+ module.conv_reg.0.mha.in_proj_weight              3145728                        0
+   module.conv_reg.0.mha.in_proj_bias                 3072                        0
+module.conv_reg.0.mha.out_proj.weight              1048576                        0
+  module.conv_reg.0.mha.out_proj.bias                 1024                        0
+       module.conv_reg.0.norm0.weight                 1024                        0
+         module.conv_reg.0.norm0.bias                 1024                        0
+       module.conv_reg.0.norm1.weight                 1024                        0
+         module.conv_reg.0.norm1.bias                 1024                        0
+       module.conv_reg.0.seq.0.weight              1048576                        0
+         module.conv_reg.0.seq.0.bias                 1024                        0
+       module.conv_reg.0.seq.2.weight              1048576                        0
+         module.conv_reg.0.seq.2.bias                 1024                        0
+ module.conv_reg.1.mha.in_proj_weight              3145728                        0
+   module.conv_reg.1.mha.in_proj_bias                 3072                        0
+module.conv_reg.1.mha.out_proj.weight              1048576                        0
+  module.conv_reg.1.mha.out_proj.bias                 1024                        0
+       module.conv_reg.1.norm0.weight                 1024                        0
+         module.conv_reg.1.norm0.bias                 1024                        0
+       module.conv_reg.1.norm1.weight                 1024                        0
+         module.conv_reg.1.norm1.bias                 1024                        0
+       module.conv_reg.1.seq.0.weight              1048576                        0
+         module.conv_reg.1.seq.0.bias                 1024                        0
+       module.conv_reg.1.seq.2.weight              1048576                        0
+         module.conv_reg.1.seq.2.bias                 1024                        0
+ module.conv_reg.2.mha.in_proj_weight              3145728                        0
+   module.conv_reg.2.mha.in_proj_bias                 3072                        0
+module.conv_reg.2.mha.out_proj.weight              1048576                        0
+  module.conv_reg.2.mha.out_proj.bias                 1024                        0
+       module.conv_reg.2.norm0.weight                 1024                        0
+         module.conv_reg.2.norm0.bias                 1024                        0
+       module.conv_reg.2.norm1.weight                 1024                        0
+         module.conv_reg.2.norm1.bias                 1024                        0
+       module.conv_reg.2.seq.0.weight              1048576                        0
+         module.conv_reg.2.seq.0.bias                 1024                        0
+       module.conv_reg.2.seq.2.weight              1048576                        0
+         module.conv_reg.2.seq.2.bias                 1024                        0
+ module.conv_reg.3.mha.in_proj_weight              3145728                        0
+   module.conv_reg.3.mha.in_proj_bias                 3072                        0
+module.conv_reg.3.mha.out_proj.weight              1048576                        0
+  module.conv_reg.3.mha.out_proj.bias                 1024                        0
+       module.conv_reg.3.norm0.weight                 1024                        0
+         module.conv_reg.3.norm0.bias                 1024                        0
+       module.conv_reg.3.norm1.weight                 1024                        0
+         module.conv_reg.3.norm1.bias                 1024                        0
+       module.conv_reg.3.seq.0.weight              1048576                        0
+         module.conv_reg.3.seq.0.bias                 1024                        0
+       module.conv_reg.3.seq.2.weight              1048576                        0
+         module.conv_reg.3.seq.2.bias                 1024                        0
+ module.conv_reg.4.mha.in_proj_weight              3145728                        0
+   module.conv_reg.4.mha.in_proj_bias                 3072                        0
+module.conv_reg.4.mha.out_proj.weight              1048576                        0
+  module.conv_reg.4.mha.out_proj.bias                 1024                        0
+       module.conv_reg.4.norm0.weight                 1024                        0
+         module.conv_reg.4.norm0.bias                 1024                        0
+       module.conv_reg.4.norm1.weight                 1024                        0
+         module.conv_reg.4.norm1.bias                 1024                        0
+       module.conv_reg.4.seq.0.weight              1048576                        0
+         module.conv_reg.4.seq.0.bias                 1024                        0
+       module.conv_reg.4.seq.2.weight              1048576                        0
+         module.conv_reg.4.seq.2.bias                 1024                        0
+ module.conv_reg.5.mha.in_proj_weight              3145728                        0
+   module.conv_reg.5.mha.in_proj_bias                 3072                        0
+module.conv_reg.5.mha.out_proj.weight              1048576                        0
+  module.conv_reg.5.mha.out_proj.bias                 1024                        0
+       module.conv_reg.5.norm0.weight                 1024                        0
+         module.conv_reg.5.norm0.bias                 1024                        0
+       module.conv_reg.5.norm1.weight                 1024                        0
+         module.conv_reg.5.norm1.bias                 1024                        0
+       module.conv_reg.5.seq.0.weight              1048576                        0
+         module.conv_reg.5.seq.0.bias                 1024                        0
+       module.conv_reg.5.seq.2.weight              1048576                        0
+         module.conv_reg.5.seq.2.bias                 1024                        0
+   module.nn_binary_particle.0.weight              1048576                        0
+     module.nn_binary_particle.0.bias                 1024                        0
+   module.nn_binary_particle.2.weight                 1024                        0
+     module.nn_binary_particle.2.bias                 1024                        0
+   module.nn_binary_particle.4.weight                 2048                        0
+     module.nn_binary_particle.4.bias                    2                        0
+               module.nn_pid.0.weight              1048576                        0
+                 module.nn_pid.0.bias                 1024                        0
+               module.nn_pid.2.weight                 1024                        0
+                 module.nn_pid.2.bias                 1024                        0
+               module.nn_pid.4.weight                 6144                        0
+                 module.nn_pid.4.bias                    6                        0
+           module.nn_pt.nn.0.0.weight              1048576                        0
+             module.nn_pt.nn.0.0.bias                 1024                        0
+           module.nn_pt.nn.0.2.weight                 1024                        0
+             module.nn_pt.nn.0.2.bias                 1024                        0
+           module.nn_pt.nn.0.4.weight                 1024                        0
+             module.nn_pt.nn.0.4.bias                    1                        0
+           module.nn_pt.nn.1.0.weight              1048576                        0
+             module.nn_pt.nn.1.0.bias                 1024                        0
+           module.nn_pt.nn.1.2.weight                 1024                        0
+             module.nn_pt.nn.1.2.bias                 1024                        0
+           module.nn_pt.nn.1.4.weight                 1024                        0
+             module.nn_pt.nn.1.4.bias                    1                        0
+            module.nn_eta.nn.0.weight              1048576                        0
+              module.nn_eta.nn.0.bias                 1024                        0
+            module.nn_eta.nn.2.weight                 1024                        0
+              module.nn_eta.nn.2.bias                 1024                        0
+            module.nn_eta.nn.4.weight                 2048                        0
+              module.nn_eta.nn.4.bias                    2                        0
+        module.nn_sin_phi.nn.0.weight              1048576                        0
+          module.nn_sin_phi.nn.0.bias                 1024                        0
+        module.nn_sin_phi.nn.2.weight                 1024                        0
+          module.nn_sin_phi.nn.2.bias                 1024                        0
+        module.nn_sin_phi.nn.4.weight                 2048                        0
+          module.nn_sin_phi.nn.4.bias                    2                        0
+        module.nn_cos_phi.nn.0.weight              1048576                        0
+          module.nn_cos_phi.nn.0.bias                 1024                        0
+        module.nn_cos_phi.nn.2.weight                 1024                        0
+          module.nn_cos_phi.nn.2.bias                 1024                        0
+        module.nn_cos_phi.nn.4.weight                 2048                        0
+          module.nn_cos_phi.nn.4.bias                    2                        0
+       module.nn_energy.nn.0.0.weight              1048576                        0
+         module.nn_energy.nn.0.0.bias                 1024                        0
+       module.nn_energy.nn.0.2.weight                 1024                        0
+         module.nn_energy.nn.0.2.bias                 1024                        0
+       module.nn_energy.nn.0.4.weight                 1024                        0
+         module.nn_energy.nn.0.4.bias                    1                        0
+       module.nn_energy.nn.1.0.weight              1048576                        0
+         module.nn_energy.nn.1.0.bias                 1024                        0
+       module.nn_energy.nn.1.2.weight                 1024                        0
+         module.nn_energy.nn.1.2.bias                 1024                        0
+       module.nn_energy.nn.1.4.weight                 1024                        0
+         module.nn_energy.nn.1.4.bias                    1                        0
+          module.final_norm_id.weight                 1024                        0
+            module.final_norm_id.bias                 1024                        0
+         module.final_norm_reg.weight                 1024                        0
+           module.final_norm_reg.bias                 1024                        0
+[2024-10-11 10:25:16,070] INFO: Creating experiment dir experiments/pyg-clic_20241011_102451_167094
+[2024-10-11 10:25:16,071] INFO: Model directory experiments/pyg-clic_20241011_102451_167094
+[2024-10-11 10:25:16,279] INFO: train_dataset: clic_edm_qq_pf, 3598296
+[2024-10-11 10:25:16,379] INFO: train_dataset: clic_edm_ttbar_pf, 7139800
+[2024-10-11 10:25:16,464] INFO: train_dataset: clic_edm_ww_fullhad_pf, 3600900
+[2024-10-11 10:25:49,796] INFO: valid_dataset: clic_edm_qq_pf, 399822
+[2024-10-11 10:25:49,821] INFO: valid_dataset: clic_edm_ttbar_pf, 793400
+[2024-10-11 10:25:49,834] INFO: valid_dataset: clic_edm_ww_fullhad_pf, 400100
+[2024-10-11 10:25:50,665] INFO: Initiating epoch #1 train run on device rank=0
+[2024-10-11 13:35:01,251] INFO: Initiating epoch #1 valid run on device rank=0
+[2024-10-11 13:42:08,318] INFO: Rank 0: epoch=1 / 30 train_loss=3.4490 valid_loss=2.8678 stale=0 epoch_train_time=189.18m epoch_valid_time=7.01m epoch_total_time=196.29m eta=5692.5m
+[2024-10-11 13:42:08,328] INFO: Initiating epoch #2 train run on device rank=0
+[2024-10-11 16:49:04,054] INFO: Initiating epoch #2 valid run on device rank=0
+[2024-10-11 16:56:07,094] INFO: Rank 0: epoch=2 / 30 train_loss=2.6138 valid_loss=2.5000 stale=0 epoch_train_time=186.93m epoch_valid_time=6.94m epoch_total_time=193.98m eta=5463.8m
+[2024-10-11 16:56:07,110] INFO: Initiating epoch #3 train run on device rank=0
+[2024-10-11 20:03:17,596] INFO: Initiating epoch #3 valid run on device rank=0
+[2024-10-11 20:10:21,873] INFO: Rank 0: epoch=3 / 30 train_loss=2.3946 valid_loss=2.3688 stale=0 epoch_train_time=187.17m epoch_valid_time=6.96m epoch_total_time=194.25m eta=5260.7m
+[2024-10-11 20:10:21,896] INFO: Initiating epoch #4 train run on device rank=0
+[2024-10-11 23:17:12,730] INFO: Initiating epoch #4 valid run on device rank=0
+[2024-10-11 23:24:13,970] INFO: Rank 0: epoch=4 / 30 train_loss=2.2850 valid_loss=2.2818 stale=0 epoch_train_time=186.85m epoch_valid_time=6.91m epoch_total_time=193.87m eta=5059.5m
+[2024-10-11 23:24:13,979] INFO: Initiating epoch #5 train run on device rank=0
+[2024-10-12 02:31:25,465] INFO: Initiating epoch #5 valid run on device rank=0
+[2024-10-12 02:38:28,338] INFO: Rank 0: epoch=5 / 30 train_loss=2.2046 valid_loss=2.2072 stale=0 epoch_train_time=187.19m epoch_valid_time=6.94m epoch_total_time=194.24m eta=4863.1m
+[2024-10-12 02:38:28,362] INFO: Initiating epoch #6 train run on device rank=0
+[2024-10-12 05:45:31,198] INFO: Initiating epoch #6 valid run on device rank=0
+[2024-10-12 05:52:32,395] INFO: Rank 0: epoch=6 / 30 train_loss=2.1378 valid_loss=2.1469 stale=0 epoch_train_time=187.05m epoch_valid_time=6.91m epoch_total_time=194.07m eta=4666.8m
+[2024-10-12 05:52:32,420] INFO: Initiating epoch #7 train run on device rank=0
+[2024-10-12 08:59:54,801] INFO: Initiating epoch #7 valid run on device rank=0
+[2024-10-12 09:06:58,357] INFO: Rank 0: epoch=7 / 30 train_loss=2.0795 valid_loss=2.0962 stale=0 epoch_train_time=187.37m epoch_valid_time=6.95m epoch_total_time=194.43m eta=4472.3m
+[2024-10-12 09:06:58,381] INFO: Initiating epoch #8 train run on device rank=0
+[2024-10-12 12:13:59,370] INFO: Initiating epoch #8 valid run on device rank=0
+[2024-10-12 12:21:02,407] INFO: Rank 0: epoch=8 / 30 train_loss=2.0358 valid_loss=2.0624 stale=0 epoch_train_time=187.02m epoch_valid_time=6.94m epoch_total_time=194.07m eta=4276.8m
+[2024-10-12 12:21:02,425] INFO: Initiating epoch #9 train run on device rank=0
+[2024-10-12 15:28:05,805] INFO: Initiating epoch #9 valid run on device rank=0
+[2024-10-12 15:35:08,750] INFO: Rank 0: epoch=9 / 30 train_loss=2.0020 valid_loss=2.0356 stale=0 epoch_train_time=187.06m epoch_valid_time=6.94m epoch_total_time=194.11m eta=4081.7m
+[2024-10-12 15:35:08,766] INFO: Initiating epoch #10 train run on device rank=0
+[2024-10-12 18:41:58,556] INFO: Initiating epoch #10 valid run on device rank=0
+[2024-10-12 18:49:01,652] INFO: Rank 0: epoch=10 / 30 train_loss=1.9755 valid_loss=2.0143 stale=0 epoch_train_time=186.83m epoch_valid_time=6.94m epoch_total_time=193.88m eta=3886.4m
+[2024-10-12 18:49:01,676] INFO: Initiating epoch #11 train run on device rank=0
+[2024-10-12 21:56:02,732] INFO: Initiating epoch #11 valid run on device rank=0
+[2024-10-12 22:03:06,330] INFO: Rank 0: epoch=11 / 30 train_loss=1.9537 valid_loss=1.9981 stale=0 epoch_train_time=187.02m epoch_valid_time=6.95m epoch_total_time=194.08m eta=3691.6m
+[2024-10-12 22:03:06,351] INFO: Initiating epoch #12 train run on device rank=0
+[2024-10-13 01:09:56,849] INFO: Initiating epoch #12 valid run on device rank=0
+[2024-10-13 01:16:58,608] INFO: Rank 0: epoch=12 / 30 train_loss=1.9351 valid_loss=1.9845 stale=0 epoch_train_time=186.84m epoch_valid_time=6.92m epoch_total_time=193.87m eta=3496.7m
+[2024-10-13 01:16:58,631] INFO: Initiating epoch #13 train run on device rank=0
+[2024-10-13 04:23:56,005] INFO: Initiating epoch #13 valid run on device rank=0
+[2024-10-13 04:30:59,653] INFO: Rank 0: epoch=13 / 30 train_loss=1.9185 valid_loss=1.9727 stale=0 epoch_train_time=186.96m epoch_valid_time=6.95m epoch_total_time=194.02m eta=3302.1m
+[2024-10-13 04:30:59,663] INFO: Initiating epoch #14 train run on device rank=0
+[2024-10-13 07:37:50,487] INFO: Initiating epoch #14 valid run on device rank=0
+[2024-10-13 07:44:53,667] INFO: Rank 0: epoch=14 / 30 train_loss=1.9034 valid_loss=1.9607 stale=0 epoch_train_time=186.85m epoch_valid_time=6.95m epoch_total_time=193.9m eta=3107.5m
+[2024-10-13 07:44:53,696] INFO: Initiating epoch #15 train run on device rank=0
+[2024-10-13 10:51:52,917] INFO: Initiating epoch #15 valid run on device rank=0
+[2024-10-13 10:58:55,298] INFO: Rank 0: epoch=15 / 30 train_loss=1.8897 valid_loss=1.9506 stale=0 epoch_train_time=186.99m epoch_valid_time=6.93m epoch_total_time=194.03m eta=2913.1m
+[2024-10-13 10:58:55,316] INFO: Initiating epoch #16 train run on device rank=0
+[2024-10-13 14:05:54,702] INFO: Initiating epoch #16 valid run on device rank=0
+[2024-10-13 14:12:57,300] INFO: Rank 0: epoch=16 / 30 train_loss=1.8772 valid_loss=1.9428 stale=0 epoch_train_time=186.99m epoch_valid_time=6.94m epoch_total_time=194.03m eta=2718.7m
+[2024-10-13 14:12:57,326] INFO: Initiating epoch #17 train run on device rank=0
+[2024-10-13 17:20:12,568] INFO: Initiating epoch #17 valid run on device rank=0
+[2024-10-13 17:27:16,171] INFO: Rank 0: epoch=17 / 30 train_loss=1.8658 valid_loss=1.9375 stale=0 epoch_train_time=187.25m epoch_valid_time=6.95m epoch_total_time=194.31m eta=2524.6m
+[2024-10-13 17:27:16,206] INFO: Initiating epoch #18 train run on device rank=0
+[2024-10-13 20:34:27,667] INFO: Initiating epoch #18 valid run on device rank=0
+[2024-10-13 20:41:31,076] INFO: Rank 0: epoch=18 / 30 train_loss=1.8551 valid_loss=1.9299 stale=0 epoch_train_time=187.19m epoch_valid_time=6.95m epoch_total_time=194.25m eta=2330.4m
+[2024-10-13 20:41:31,089] INFO: Initiating epoch #19 train run on device rank=0
+[2024-10-13 23:48:44,511] INFO: Initiating epoch #19 valid run on device rank=0
+[2024-10-13 23:55:47,205] INFO: Rank 0: epoch=19 / 30 train_loss=1.8452 valid_loss=1.9244 stale=0 epoch_train_time=187.22m epoch_valid_time=6.94m epoch_total_time=194.27m eta=2136.3m
+[2024-10-13 23:55:47,222] INFO: Initiating epoch #20 train run on device rank=0
+[2024-10-14 03:02:39,426] INFO: Initiating epoch #20 valid run on device rank=0
+[2024-10-14 03:09:42,527] INFO: Rank 0: epoch=20 / 30 train_loss=1.8362 valid_loss=1.9202 stale=0 epoch_train_time=186.87m epoch_valid_time=6.94m epoch_total_time=193.92m eta=1941.9m
+[2024-10-14 03:09:42,550] INFO: Initiating epoch #21 train run on device rank=0
+[2024-10-14 06:16:45,945] INFO: Initiating epoch #21 valid run on device rank=0
+[2024-10-14 06:23:49,200] INFO: Rank 0: epoch=21 / 30 train_loss=1.8278 valid_loss=1.9160 stale=0 epoch_train_time=187.06m epoch_valid_time=6.95m epoch_total_time=194.11m eta=1747.7m
+[2024-10-14 06:23:49,218] INFO: Initiating epoch #22 train run on device rank=0
+[2024-10-14 09:30:41,408] INFO: Initiating epoch #22 valid run on device rank=0
+[2024-10-14 09:37:46,117] INFO: Rank 0: epoch=22 / 30 train_loss=1.8199 valid_loss=1.9131 stale=0 epoch_train_time=186.87m epoch_valid_time=6.97m epoch_total_time=193.95m eta=1553.4m
+[2024-10-14 09:37:46,143] INFO: Initiating epoch #23 train run on device rank=0
+[2024-10-14 09:47:43,648] INFO: Will use torch.nn.parallel.DistributedDataParallel() and 8 gpus
+[2024-10-14 09:47:43,652] INFO: AMD Radeon Graphics
+[2024-10-14 09:47:43,652] INFO: AMD Radeon Graphics
+[2024-10-14 09:47:43,652] INFO: AMD Radeon Graphics
+[2024-10-14 09:47:43,652] INFO: AMD Radeon Graphics
+[2024-10-14 09:47:43,652] INFO: AMD Radeon Graphics
+[2024-10-14 09:47:43,652] INFO: AMD Radeon Graphics
+[2024-10-14 09:47:43,653] INFO: AMD Radeon Graphics
+[2024-10-14 09:47:43,653] INFO: AMD Radeon Graphics
+[2024-10-14 09:47:47,850] INFO: configured dtype=torch.bfloat16 for autocast
+[2024-10-14 09:47:50,639] INFO: model_kwargs: {'input_dim': 17, 'num_classes': 6, 'input_encoding': 'split', 'pt_mode': 'direct-elemtype-split', 'eta_mode': 'linear', 'sin_phi_mode': 'linear', 'cos_phi_mode': 'linear', 'energy_mode': 'direct-elemtype-split', 'elemtypes_nonzero': [1, 2], 'learned_representation_mode': 'last', 'conv_type': 'attention', 'num_convs': 6, 'dropout_ff': 0.1, 'dropout_conv_id_mha': 0.0, 'dropout_conv_id_ff': 0.0, 'dropout_conv_reg_mha': 0.1, 'dropout_conv_reg_ff': 0.1, 'activation': 'gelu', 'head_dim': 32, 'num_heads': 32, 'attention_type': 'math', 'use_pre_layernorm': True}
+[2024-10-14 09:47:50,738] INFO: using attention_type=math
+[2024-10-14 09:47:50,778] INFO: using attention_type=math
+[2024-10-14 09:47:50,814] INFO: using attention_type=math
+[2024-10-14 09:47:50,849] INFO: using attention_type=math
+[2024-10-14 09:47:50,884] INFO: using attention_type=math
+[2024-10-14 09:47:50,921] INFO: using attention_type=math
+[2024-10-14 09:47:50,957] INFO: using attention_type=math
+[2024-10-14 09:47:50,993] INFO: using attention_type=math
+[2024-10-14 09:47:51,029] INFO: using attention_type=math
+[2024-10-14 09:47:51,065] INFO: using attention_type=math
+[2024-10-14 09:47:51,102] INFO: using attention_type=math
+[2024-10-14 09:47:51,137] INFO: using attention_type=math
+[2024-10-14 09:47:58,146] INFO: Loaded model weights from experiments/pyg-clic_20241011_102451_167094/checkpoints/checkpoint-22-1.913142.pth
+[2024-10-14 09:47:59,962] INFO: DistributedDataParallel(
+  (module): MLPF(
+    (nn0_id): ModuleList(
+      (0-1): 2 x Sequential(
+        (0): Linear(in_features=17, out_features=1024, bias=True)
+        (1): GELU(approximate='none')
+        (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (3): Dropout(p=0.1, inplace=False)
+        (4): Linear(in_features=1024, out_features=1024, bias=True)
+      )
+    )
+    (nn0_reg): ModuleList(
+      (0-1): 2 x Sequential(
+        (0): Linear(in_features=17, out_features=1024, bias=True)
+        (1): GELU(approximate='none')
+        (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (3): Dropout(p=0.1, inplace=False)
+        (4): Linear(in_features=1024, out_features=1024, bias=True)
+      )
+    )
+    (conv_id): ModuleList(
+      (0-5): 6 x PreLnSelfAttentionLayer(
+        (mha): MultiheadAttention(
+          (out_proj): NonDynamicallyQuantizableLinear(in_features=1024, out_features=1024, bias=True)
+        )
+        (norm0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (seq): Sequential(
+          (0): Linear(in_features=1024, out_features=1024, bias=True)
+          (1): GELU(approximate='none')
+          (2): Linear(in_features=1024, out_features=1024, bias=True)
+          (3): GELU(approximate='none')
+        )
+        (dropout): Dropout(p=0.0, inplace=False)
+      )
+    )
+    (conv_reg): ModuleList(
+      (0-5): 6 x PreLnSelfAttentionLayer(
+        (mha): MultiheadAttention(
+          (out_proj): NonDynamicallyQuantizableLinear(in_features=1024, out_features=1024, bias=True)
+        )
+        (norm0): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (seq): Sequential(
+          (0): Linear(in_features=1024, out_features=1024, bias=True)
+          (1): GELU(approximate='none')
+          (2): Linear(in_features=1024, out_features=1024, bias=True)
+          (3): GELU(approximate='none')
+        )
+        (dropout): Dropout(p=0.1, inplace=False)
+      )
+    )
+    (nn_binary_particle): Sequential(
+      (0): Linear(in_features=1024, out_features=1024, bias=True)
+      (1): GELU(approximate='none')
+      (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+      (3): Dropout(p=0.1, inplace=False)
+      (4): Linear(in_features=1024, out_features=2, bias=True)
+    )
+    (nn_pid): Sequential(
+      (0): Linear(in_features=1024, out_features=1024, bias=True)
+      (1): GELU(approximate='none')
+      (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+      (3): Dropout(p=0.1, inplace=False)
+      (4): Linear(in_features=1024, out_features=6, bias=True)
+    )
+    (nn_pt): RegressionOutput(
+      (nn): ModuleList(
+        (0-1): 2 x Sequential(
+          (0): Linear(in_features=1024, out_features=1024, bias=True)
+          (1): GELU(approximate='none')
+          (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+          (3): Dropout(p=0.1, inplace=False)
+          (4): Linear(in_features=1024, out_features=1, bias=True)
+        )
+      )
+    )
+    (nn_eta): RegressionOutput(
+      (nn): Sequential(
+        (0): Linear(in_features=1024, out_features=1024, bias=True)
+        (1): GELU(approximate='none')
+        (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (3): Dropout(p=0.1, inplace=False)
+        (4): Linear(in_features=1024, out_features=2, bias=True)
+      )
+    )
+    (nn_sin_phi): RegressionOutput(
+      (nn): Sequential(
+        (0): Linear(in_features=1024, out_features=1024, bias=True)
+        (1): GELU(approximate='none')
+        (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (3): Dropout(p=0.1, inplace=False)
+        (4): Linear(in_features=1024, out_features=2, bias=True)
+      )
+    )
+    (nn_cos_phi): RegressionOutput(
+      (nn): Sequential(
+        (0): Linear(in_features=1024, out_features=1024, bias=True)
+        (1): GELU(approximate='none')
+        (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+        (3): Dropout(p=0.1, inplace=False)
+        (4): Linear(in_features=1024, out_features=2, bias=True)
+      )
+    )
+    (nn_energy): RegressionOutput(
+      (nn): ModuleList(
+        (0-1): 2 x Sequential(
+          (0): Linear(in_features=1024, out_features=1024, bias=True)
+          (1): GELU(approximate='none')
+          (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+          (3): Dropout(p=0.1, inplace=False)
+          (4): Linear(in_features=1024, out_features=1, bias=True)
+        )
+      )
+    )
+    (final_norm_id): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+    (final_norm_reg): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
+  )
+)
+[2024-10-14 09:47:59,965] INFO: Trainable parameters: 89388050
+[2024-10-14 09:47:59,965] INFO: Non-trainable parameters: 0
+[2024-10-14 09:47:59,965] INFO: Total parameters: 89388050
+[2024-10-14 09:47:59,972] INFO:                               Modules Trainable parameters Non-trainable parameters
+             module.nn0_id.0.0.weight                17408                        0
+               module.nn0_id.0.0.bias                 1024                        0
+             module.nn0_id.0.2.weight                 1024                        0
+               module.nn0_id.0.2.bias                 1024                        0
+             module.nn0_id.0.4.weight              1048576                        0
+               module.nn0_id.0.4.bias                 1024                        0
+             module.nn0_id.1.0.weight                17408                        0
+               module.nn0_id.1.0.bias                 1024                        0
+             module.nn0_id.1.2.weight                 1024                        0
+               module.nn0_id.1.2.bias                 1024                        0
+             module.nn0_id.1.4.weight              1048576                        0
+               module.nn0_id.1.4.bias                 1024                        0
+            module.nn0_reg.0.0.weight                17408                        0
+              module.nn0_reg.0.0.bias                 1024                        0
+            module.nn0_reg.0.2.weight                 1024                        0
+              module.nn0_reg.0.2.bias                 1024                        0
+            module.nn0_reg.0.4.weight              1048576                        0
+              module.nn0_reg.0.4.bias                 1024                        0
+            module.nn0_reg.1.0.weight                17408                        0
+              module.nn0_reg.1.0.bias                 1024                        0
+            module.nn0_reg.1.2.weight                 1024                        0
+              module.nn0_reg.1.2.bias                 1024                        0
+            module.nn0_reg.1.4.weight              1048576                        0
+              module.nn0_reg.1.4.bias                 1024                        0
+  module.conv_id.0.mha.in_proj_weight              3145728                        0
+    module.conv_id.0.mha.in_proj_bias                 3072                        0
+ module.conv_id.0.mha.out_proj.weight              1048576                        0
+   module.conv_id.0.mha.out_proj.bias                 1024                        0
+        module.conv_id.0.norm0.weight                 1024                        0
+          module.conv_id.0.norm0.bias                 1024                        0
+        module.conv_id.0.norm1.weight                 1024                        0
+          module.conv_id.0.norm1.bias                 1024                        0
+        module.conv_id.0.seq.0.weight              1048576                        0
+          module.conv_id.0.seq.0.bias                 1024                        0
+        module.conv_id.0.seq.2.weight              1048576                        0
+          module.conv_id.0.seq.2.bias                 1024                        0
+  module.conv_id.1.mha.in_proj_weight              3145728                        0
+    module.conv_id.1.mha.in_proj_bias                 3072                        0
+ module.conv_id.1.mha.out_proj.weight              1048576                        0
+   module.conv_id.1.mha.out_proj.bias                 1024                        0
+        module.conv_id.1.norm0.weight                 1024                        0
+          module.conv_id.1.norm0.bias                 1024                        0
+        module.conv_id.1.norm1.weight                 1024                        0
+          module.conv_id.1.norm1.bias                 1024                        0
+        module.conv_id.1.seq.0.weight              1048576                        0
+          module.conv_id.1.seq.0.bias                 1024                        0
+        module.conv_id.1.seq.2.weight              1048576                        0
+          module.conv_id.1.seq.2.bias                 1024                        0
+  module.conv_id.2.mha.in_proj_weight              3145728                        0
+    module.conv_id.2.mha.in_proj_bias                 3072                        0
+ module.conv_id.2.mha.out_proj.weight              1048576                        0
+   module.conv_id.2.mha.out_proj.bias                 1024                        0
+        module.conv_id.2.norm0.weight                 1024                        0
+          module.conv_id.2.norm0.bias                 1024                        0
+        module.conv_id.2.norm1.weight                 1024                        0
+          module.conv_id.2.norm1.bias                 1024                        0
+        module.conv_id.2.seq.0.weight              1048576                        0
+          module.conv_id.2.seq.0.bias                 1024                        0
+        module.conv_id.2.seq.2.weight              1048576                        0
+          module.conv_id.2.seq.2.bias                 1024                        0
+  module.conv_id.3.mha.in_proj_weight              3145728                        0
+    module.conv_id.3.mha.in_proj_bias                 3072                        0
+ module.conv_id.3.mha.out_proj.weight              1048576                        0
+   module.conv_id.3.mha.out_proj.bias                 1024                        0
+        module.conv_id.3.norm0.weight                 1024                        0
+          module.conv_id.3.norm0.bias                 1024                        0
+        module.conv_id.3.norm1.weight                 1024                        0
+          module.conv_id.3.norm1.bias                 1024                        0
+        module.conv_id.3.seq.0.weight              1048576                        0
+          module.conv_id.3.seq.0.bias                 1024                        0
+        module.conv_id.3.seq.2.weight              1048576                        0
+          module.conv_id.3.seq.2.bias                 1024                        0
+  module.conv_id.4.mha.in_proj_weight              3145728                        0
+    module.conv_id.4.mha.in_proj_bias                 3072                        0
+ module.conv_id.4.mha.out_proj.weight              1048576                        0
+   module.conv_id.4.mha.out_proj.bias                 1024                        0
+        module.conv_id.4.norm0.weight                 1024                        0
+          module.conv_id.4.norm0.bias                 1024                        0
+        module.conv_id.4.norm1.weight                 1024                        0
+          module.conv_id.4.norm1.bias                 1024                        0
+        module.conv_id.4.seq.0.weight              1048576                        0
+          module.conv_id.4.seq.0.bias                 1024                        0
+        module.conv_id.4.seq.2.weight              1048576                        0
+          module.conv_id.4.seq.2.bias                 1024                        0
+  module.conv_id.5.mha.in_proj_weight              3145728                        0
+    module.conv_id.5.mha.in_proj_bias                 3072                        0
+ module.conv_id.5.mha.out_proj.weight              1048576                        0
+   module.conv_id.5.mha.out_proj.bias                 1024                        0
+        module.conv_id.5.norm0.weight                 1024                        0
+          module.conv_id.5.norm0.bias                 1024                        0
+        module.conv_id.5.norm1.weight                 1024                        0
+          module.conv_id.5.norm1.bias                 1024                        0
+        module.conv_id.5.seq.0.weight              1048576                        0
+          module.conv_id.5.seq.0.bias                 1024                        0
+        module.conv_id.5.seq.2.weight              1048576                        0
+          module.conv_id.5.seq.2.bias                 1024                        0
+ module.conv_reg.0.mha.in_proj_weight              3145728                        0
+   module.conv_reg.0.mha.in_proj_bias                 3072                        0
+module.conv_reg.0.mha.out_proj.weight              1048576                        0
+  module.conv_reg.0.mha.out_proj.bias                 1024                        0
+       module.conv_reg.0.norm0.weight                 1024                        0
+         module.conv_reg.0.norm0.bias                 1024                        0
+       module.conv_reg.0.norm1.weight                 1024                        0
+         module.conv_reg.0.norm1.bias                 1024                        0
+       module.conv_reg.0.seq.0.weight              1048576                        0
+         module.conv_reg.0.seq.0.bias                 1024                        0
+       module.conv_reg.0.seq.2.weight              1048576                        0
+         module.conv_reg.0.seq.2.bias                 1024                        0
+ module.conv_reg.1.mha.in_proj_weight              3145728                        0
+   module.conv_reg.1.mha.in_proj_bias                 3072                        0
+module.conv_reg.1.mha.out_proj.weight              1048576                        0
+  module.conv_reg.1.mha.out_proj.bias                 1024                        0
+       module.conv_reg.1.norm0.weight                 1024                        0
+         module.conv_reg.1.norm0.bias                 1024                        0
+       module.conv_reg.1.norm1.weight                 1024                        0
+         module.conv_reg.1.norm1.bias                 1024                        0
+       module.conv_reg.1.seq.0.weight              1048576                        0
+         module.conv_reg.1.seq.0.bias                 1024                        0
+       module.conv_reg.1.seq.2.weight              1048576                        0
+         module.conv_reg.1.seq.2.bias                 1024                        0
+ module.conv_reg.2.mha.in_proj_weight              3145728                        0
+   module.conv_reg.2.mha.in_proj_bias                 3072                        0
+module.conv_reg.2.mha.out_proj.weight              1048576                        0
+  module.conv_reg.2.mha.out_proj.bias                 1024                        0
+       module.conv_reg.2.norm0.weight                 1024                        0
+         module.conv_reg.2.norm0.bias                 1024                        0
+       module.conv_reg.2.norm1.weight                 1024                        0
+         module.conv_reg.2.norm1.bias                 1024                        0
+       module.conv_reg.2.seq.0.weight              1048576                        0
+         module.conv_reg.2.seq.0.bias                 1024                        0
+       module.conv_reg.2.seq.2.weight              1048576                        0
+         module.conv_reg.2.seq.2.bias                 1024                        0
+ module.conv_reg.3.mha.in_proj_weight              3145728                        0
+   module.conv_reg.3.mha.in_proj_bias                 3072                        0
+module.conv_reg.3.mha.out_proj.weight              1048576                        0
+  module.conv_reg.3.mha.out_proj.bias                 1024                        0
+       module.conv_reg.3.norm0.weight                 1024                        0
+         module.conv_reg.3.norm0.bias                 1024                        0
+       module.conv_reg.3.norm1.weight                 1024                        0
+         module.conv_reg.3.norm1.bias                 1024                        0
+       module.conv_reg.3.seq.0.weight              1048576                        0
+         module.conv_reg.3.seq.0.bias                 1024                        0
+       module.conv_reg.3.seq.2.weight              1048576                        0
+         module.conv_reg.3.seq.2.bias                 1024                        0
+ module.conv_reg.4.mha.in_proj_weight              3145728                        0
+   module.conv_reg.4.mha.in_proj_bias                 3072                        0
+module.conv_reg.4.mha.out_proj.weight              1048576                        0
+  module.conv_reg.4.mha.out_proj.bias                 1024                        0
+       module.conv_reg.4.norm0.weight                 1024                        0
+         module.conv_reg.4.norm0.bias                 1024                        0
+       module.conv_reg.4.norm1.weight                 1024                        0
+         module.conv_reg.4.norm1.bias                 1024                        0
+       module.conv_reg.4.seq.0.weight              1048576                        0
+         module.conv_reg.4.seq.0.bias                 1024                        0
+       module.conv_reg.4.seq.2.weight              1048576                        0
+         module.conv_reg.4.seq.2.bias                 1024                        0
+ module.conv_reg.5.mha.in_proj_weight              3145728                        0
+   module.conv_reg.5.mha.in_proj_bias                 3072                        0
+module.conv_reg.5.mha.out_proj.weight              1048576                        0
+  module.conv_reg.5.mha.out_proj.bias                 1024                        0
+       module.conv_reg.5.norm0.weight                 1024                        0
+         module.conv_reg.5.norm0.bias                 1024                        0
+       module.conv_reg.5.norm1.weight                 1024                        0
+         module.conv_reg.5.norm1.bias                 1024                        0
+       module.conv_reg.5.seq.0.weight              1048576                        0
+         module.conv_reg.5.seq.0.bias                 1024                        0
+       module.conv_reg.5.seq.2.weight              1048576                        0
+         module.conv_reg.5.seq.2.bias                 1024                        0
+   module.nn_binary_particle.0.weight              1048576                        0
+     module.nn_binary_particle.0.bias                 1024                        0
+   module.nn_binary_particle.2.weight                 1024                        0
+     module.nn_binary_particle.2.bias                 1024                        0
+   module.nn_binary_particle.4.weight                 2048                        0
+     module.nn_binary_particle.4.bias                    2                        0
+               module.nn_pid.0.weight              1048576                        0
+                 module.nn_pid.0.bias                 1024                        0
+               module.nn_pid.2.weight                 1024                        0
+                 module.nn_pid.2.bias                 1024                        0
+               module.nn_pid.4.weight                 6144                        0
+                 module.nn_pid.4.bias                    6                        0
+           module.nn_pt.nn.0.0.weight              1048576                        0
+             module.nn_pt.nn.0.0.bias                 1024                        0
+           module.nn_pt.nn.0.2.weight                 1024                        0
+             module.nn_pt.nn.0.2.bias                 1024                        0
+           module.nn_pt.nn.0.4.weight                 1024                        0
+             module.nn_pt.nn.0.4.bias                    1                        0
+           module.nn_pt.nn.1.0.weight              1048576                        0
+             module.nn_pt.nn.1.0.bias                 1024                        0
+           module.nn_pt.nn.1.2.weight                 1024                        0
+             module.nn_pt.nn.1.2.bias                 1024                        0
+           module.nn_pt.nn.1.4.weight                 1024                        0
+             module.nn_pt.nn.1.4.bias                    1                        0
+            module.nn_eta.nn.0.weight              1048576                        0
+              module.nn_eta.nn.0.bias                 1024                        0
+            module.nn_eta.nn.2.weight                 1024                        0
+              module.nn_eta.nn.2.bias                 1024                        0
+            module.nn_eta.nn.4.weight                 2048                        0
+              module.nn_eta.nn.4.bias                    2                        0
+        module.nn_sin_phi.nn.0.weight              1048576                        0
+          module.nn_sin_phi.nn.0.bias                 1024                        0
+        module.nn_sin_phi.nn.2.weight                 1024                        0
+          module.nn_sin_phi.nn.2.bias                 1024                        0
+        module.nn_sin_phi.nn.4.weight                 2048                        0
+          module.nn_sin_phi.nn.4.bias                    2                        0
+        module.nn_cos_phi.nn.0.weight              1048576                        0
+          module.nn_cos_phi.nn.0.bias                 1024                        0
+        module.nn_cos_phi.nn.2.weight                 1024                        0
+          module.nn_cos_phi.nn.2.bias                 1024                        0
+        module.nn_cos_phi.nn.4.weight                 2048                        0
+          module.nn_cos_phi.nn.4.bias                    2                        0
+       module.nn_energy.nn.0.0.weight              1048576                        0
+         module.nn_energy.nn.0.0.bias                 1024                        0
+       module.nn_energy.nn.0.2.weight                 1024                        0
+         module.nn_energy.nn.0.2.bias                 1024                        0
+       module.nn_energy.nn.0.4.weight                 1024                        0
+         module.nn_energy.nn.0.4.bias                    1                        0
+       module.nn_energy.nn.1.0.weight              1048576                        0
+         module.nn_energy.nn.1.0.bias                 1024                        0
+       module.nn_energy.nn.1.2.weight                 1024                        0
+         module.nn_energy.nn.1.2.bias                 1024                        0
+       module.nn_energy.nn.1.4.weight                 1024                        0
+         module.nn_energy.nn.1.4.bias                    1                        0
+          module.final_norm_id.weight                 1024                        0
+            module.final_norm_id.bias                 1024                        0
+         module.final_norm_reg.weight                 1024                        0
+           module.final_norm_reg.bias                 1024                        0
+[2024-10-14 09:47:59,977] INFO: Creating experiment dir experiments/pyg-clic_20241011_102451_167094
+[2024-10-14 09:47:59,978] INFO: Model directory experiments/pyg-clic_20241011_102451_167094
+[2024-10-14 09:48:00,140] INFO: train_dataset: clic_edm_qq_pf, 3598296
+[2024-10-14 09:48:00,234] INFO: train_dataset: clic_edm_ttbar_pf, 7139800
+[2024-10-14 09:48:00,332] INFO: train_dataset: clic_edm_ww_fullhad_pf, 3600900
+[2024-10-14 09:48:35,829] INFO: valid_dataset: clic_edm_qq_pf, 399822
+[2024-10-14 09:48:35,938] INFO: valid_dataset: clic_edm_ttbar_pf, 793400
+[2024-10-14 09:48:35,965] INFO: valid_dataset: clic_edm_ww_fullhad_pf, 400100
+[2024-10-14 09:48:46,758] INFO: Initiating epoch #23 train run on device rank=0
+[2024-10-14 13:00:19,246] INFO: Initiating epoch #23 valid run on device rank=0
+[2024-10-14 13:07:20,574] INFO: Rank 0: epoch=23 / 30 train_loss=1.8128 valid_loss=1.9104 stale=0 epoch_train_time=191.54m epoch_valid_time=6.91m epoch_total_time=198.56m eta=60.4m
+[2024-10-14 13:07:20,598] INFO: Initiating epoch #24 train run on device rank=0
+[2024-10-14 16:18:48,042] INFO: Initiating epoch #24 valid run on device rank=0
+[2024-10-14 16:25:50,924] INFO: Rank 0: epoch=24 / 30 train_loss=1.8063 valid_loss=1.9085 stale=0 epoch_train_time=191.46m epoch_valid_time=6.93m epoch_total_time=198.51m eta=99.3m
+[2024-10-14 16:25:50,944] INFO: Initiating epoch #25 train run on device rank=0
+[2024-10-14 19:37:17,620] INFO: Initiating epoch #25 valid run on device rank=0
+[2024-10-14 19:44:18,373] INFO: Rank 0: epoch=25 / 30 train_loss=1.8005 valid_loss=1.9052 stale=0 epoch_train_time=191.44m epoch_valid_time=6.91m epoch_total_time=198.46m eta=119.1m
+[2024-10-14 19:44:18,398] INFO: Initiating epoch #26 train run on device rank=0
+[2024-10-14 22:55:38,092] INFO: Initiating epoch #26 valid run on device rank=0
+[2024-10-14 23:02:41,805] INFO: Rank 0: epoch=26 / 30 train_loss=1.7954 valid_loss=1.9043 stale=0 epoch_train_time=191.33m epoch_valid_time=6.95m epoch_total_time=198.39m eta=122.1m
+[2024-10-14 23:02:41,835] INFO: Initiating epoch #27 train run on device rank=0
+[2024-10-15 02:14:30,888] INFO: Initiating epoch #27 valid run on device rank=0
+[2024-10-15 02:21:32,802] INFO: Rank 0: epoch=27 / 30 train_loss=1.7909 valid_loss=1.9021 stale=0 epoch_train_time=191.82m epoch_valid_time=6.92m epoch_total_time=198.85m eta=110.3m
+[2024-10-15 02:21:32,828] INFO: Initiating epoch #28 train run on device rank=0
+[2024-10-15 05:32:29,453] INFO: Initiating epoch #28 valid run on device rank=0
+[2024-10-15 05:39:31,630] INFO: Rank 0: epoch=28 / 30 train_loss=1.7872 valid_loss=1.9020 stale=0 epoch_train_time=190.94m epoch_valid_time=6.92m epoch_total_time=197.98m eta=85.1m
+[2024-10-15 05:39:31,652] INFO: Initiating epoch #29 train run on device rank=0
+[2024-10-15 08:51:11,754] INFO: Initiating epoch #29 valid run on device rank=0
+[2024-10-15 08:58:11,804] INFO: Rank 0: epoch=29 / 30 train_loss=1.7842 valid_loss=1.9017 stale=0 epoch_train_time=191.67m epoch_valid_time=6.89m epoch_total_time=198.67m eta=47.9m
+[2024-10-15 08:58:11,830] INFO: Initiating epoch #30 train run on device rank=0
+[2024-10-15 12:09:34,952] INFO: Initiating epoch #30 valid run on device rank=0