diff --git a/Makefile b/Makefile
index 185712f..d853692 100644
--- a/Makefile
+++ b/Makefile
@@ -130,6 +130,12 @@ finetune: requirements
 finetune_cv: requirements
 	$(PYTHON_INTERPRETER) -m lnp_ml.modeling.train_cv main --init-from-pretrain models/pretrain_delivery.pt $(FREEZE_FLAG) $(MPNN_FLAG) $(DEVICE_FLAG)
 
+## Train with cross-validation on internal data only (5-fold, amine-based split)
+.PHONY: train_cv
+train_cv: requirements
+	$(PYTHON_INTERPRETER) -m lnp_ml.modeling.train_cv main $(FREEZE_FLAG) $(MPNN_FLAG) $(DEVICE_FLAG)
+
+
 ## Evaluate CV finetuned models on test sets (auto-detects MPNN from checkpoint)
 .PHONY: test_cv
 test_cv: requirements
diff --git a/models/finetune_cv/config.json b/models/finetune_cv/config.json
index ac79f18..4f7fad2 100644
--- a/models/finetune_cv/config.json
+++ b/models/finetune_cv/config.json
@@ -11,6 +11,6 @@
   "batch_size": 32,
   "epochs": 100,
   "patience": 15,
-  "init_from_pretrain": "models/pretrain_delivery.pt",
+  "init_from_pretrain": null,
   "freeze_backbone": false
 }
\ No newline at end of file
diff --git a/models/finetune_cv/cv_results.json b/models/finetune_cv/cv_results.json
index ebedc3e..a0e1664 100644
--- a/models/finetune_cv/cv_results.json
+++ b/models/finetune_cv/cv_results.json
@@ -2,38 +2,38 @@
   "fold_results": [
     {
       "fold_idx": 0,
-      "best_val_loss": 6.144314289093018,
-      "epochs_trained": 25,
-      "final_train_loss": 1.4692220211029052
+      "best_val_loss": 5.7676777839660645,
+      "epochs_trained": 24,
+      "final_train_loss": 1.4942118644714355
     },
     {
       "fold_idx": 1,
-      "best_val_loss": 8.569346030553183,
+      "best_val_loss": 8.418675899505615,
       "epochs_trained": 20,
-      "final_train_loss": 1.5929443359375
+      "final_train_loss": 1.4902493238449097
     },
     {
       "fold_idx": 2,
-      "best_val_loss": 3.7409281730651855,
-      "epochs_trained": 22,
-      "final_train_loss": 1.9401288827260335
+      "best_val_loss": 3.5122547830854143,
+      "epochs_trained": 25,
+      "final_train_loss": 1.7609570423762004
     },
     {
       "fold_idx": 3,
-      "best_val_loss": 3.47284197807312,
-      "epochs_trained": 27,
-      "final_train_loss": 1.8295514345169068
+      "best_val_loss": 3.165306806564331,
+      "epochs_trained": 21,
+      "final_train_loss": 2.0073827385902403
     },
     {
       "fold_idx": 4,
-      "best_val_loss": 2.756531000137329,
-      "epochs_trained": 19,
-      "final_train_loss": 1.9399811571294612
+      "best_val_loss": 2.996154228846232,
+      "epochs_trained": 18,
+      "final_train_loss": 1.9732873006300493
     }
   ],
   "summary": {
-    "val_loss_mean": 4.936792294184367,
-    "val_loss_std": 2.1438440638412697
+    "val_loss_mean": 4.772013900393532,
+    "val_loss_std": 2.0790222989111475
   },
   "config": {
     "d_model": 256,
@@ -48,7 +48,7 @@
     "batch_size": 32,
     "epochs": 100,
     "patience": 15,
-    "init_from_pretrain": "models/pretrain_delivery.pt",
+    "init_from_pretrain": null,
     "freeze_backbone": false
   }
 }
\ No newline at end of file
diff --git a/models/finetune_cv/fold_0/history.json b/models/finetune_cv/fold_0/history.json
index 6b58a67..8f97876 100644
--- a/models/finetune_cv/fold_0/history.json
+++ b/models/finetune_cv/fold_0/history.json
@@ -1,530 +1,509 @@
 {
   "train": [
     {
-      "loss": 19.238220310211183,
-      "loss_size": 14.334759521484376,
-      "loss_pdi": 1.275341796875,
-      "loss_ee": 1.078886091709137,
-      "loss_delivery": 0.639056247472763,
-      "loss_biodist": 1.314099133014679,
-      "loss_toxic": 0.5960778951644897
+      "loss": 17.65872812271118,
+      "loss_size": 12.601411867141724,
+      "loss_pdi": 1.3666706204414367,
+      "loss_ee": 1.0830313920974732,
+      "loss_delivery": 0.5962779104709626,
+      "loss_biodist": 1.3918164849281311,
+      "loss_toxic": 0.6195200622081757
     },
     {
-      "loss": 7.8008105754852295,
-      "loss_size": 3.835961139202118,
-      "loss_pdi": 1.0304630517959594,
-      "loss_ee": 1.002296370267868,
-      "loss_delivery": 0.527982234954834,
-      "loss_biodist": 1.0791441202163696,
-      "loss_toxic": 0.3249638095498085
+      "loss": 5.925264883041382,
+      "loss_size": 1.8580878481268883,
+      "loss_pdi": 1.1011681258678436,
+      "loss_ee": 0.971046245098114,
+      "loss_delivery": 0.5075224950909615,
+      "loss_biodist": 1.1051940202713013,
+      "loss_toxic": 0.38224617540836336
     },
     {
-      "loss": 3.952784705162048,
-      "loss_size": 0.5930101618170738,
-      "loss_pdi": 0.7961886763572693,
-      "loss_ee": 0.9416749358177186,
-      "loss_delivery": 0.5073600560426712,
-      "loss_biodist": 0.9171513140201568,
-      "loss_toxic": 0.19739954844117164
+      "loss": 3.4781792640686033,
+      "loss_size": 0.23610344529151917,
+      "loss_pdi": 0.8137399554252625,
+      "loss_ee": 0.9135127127170563,
+      "loss_delivery": 0.4596045270562172,
+      "loss_biodist": 0.8695587992668152,
+      "loss_toxic": 0.18565986081957817
     },
     {
-      "loss": 3.218132185935974,
-      "loss_size": 0.20842453986406326,
-      "loss_pdi": 0.688220864534378,
-      "loss_ee": 0.904784232378006,
-      "loss_delivery": 0.4910900041460991,
-      "loss_biodist": 0.792213362455368,
-      "loss_toxic": 0.1333992186933756
+      "loss": 2.9488561868667604,
+      "loss_size": 0.23130029290914536,
+      "loss_pdi": 0.644479614496231,
+      "loss_ee": 0.8721524059772492,
+      "loss_delivery": 0.4146773874759674,
+      "loss_biodist": 0.646893310546875,
+      "loss_toxic": 0.13935319259762763
     },
     {
-      "loss": 2.930907416343689,
-      "loss_size": 0.21291286423802375,
-      "loss_pdi": 0.6122969090938568,
-      "loss_ee": 0.8774014472961426,
-      "loss_delivery": 0.4451231583952904,
-      "loss_biodist": 0.6634868443012237,
-      "loss_toxic": 0.11968618221580982
+      "loss": 2.6432241678237913,
+      "loss_size": 0.16843259893357754,
+      "loss_pdi": 0.5857123643159866,
+      "loss_ee": 0.8315786123275757,
+      "loss_delivery": 0.4049036353826523,
+      "loss_biodist": 0.5410242855548859,
+      "loss_toxic": 0.11157271154224872
     },
     {
-      "loss": 2.7193881273269653,
-      "loss_size": 0.213371854275465,
-      "loss_pdi": 0.5864351749420166,
-      "loss_ee": 0.8563571333885193,
-      "loss_delivery": 0.3963193610310555,
-      "loss_biodist": 0.5644777715206146,
-      "loss_toxic": 0.10242685079574584
+      "loss": 2.461507487297058,
+      "loss_size": 0.18602822050452233,
+      "loss_pdi": 0.5872043997049332,
+      "loss_ee": 0.8179578661918641,
+      "loss_delivery": 0.32779163047671317,
+      "loss_biodist": 0.45097417533397677,
+      "loss_toxic": 0.09155115596950054
     },
     {
-      "loss": 2.5172106266021728,
-      "loss_size": 0.23241330087184905,
-      "loss_pdi": 0.5564007371664047,
-      "loss_ee": 0.8135433554649353,
-      "loss_delivery": 0.39191135168075564,
-      "loss_biodist": 0.4503189116716385,
-      "loss_toxic": 0.07262293715029955
+      "loss": 2.3792370796203612,
+      "loss_size": 0.2090120367705822,
+      "loss_pdi": 0.5358257800340652,
+      "loss_ee": 0.8088949501514435,
+      "loss_delivery": 0.3434994474053383,
+      "loss_biodist": 0.40993946194648745,
+      "loss_toxic": 0.07206540685147048
     },
     {
-      "loss": 2.3014568209648134,
-      "loss_size": 0.1924597330391407,
-      "loss_pdi": 0.5315678030252456,
-      "loss_ee": 0.8107137799263,
-      "loss_delivery": 0.33130097687244414,
-      "loss_biodist": 0.3714979439973831,
-      "loss_toxic": 0.06391655802726745
+      "loss": 2.207099366188049,
+      "loss_size": 0.1589151345193386,
+      "loss_pdi": 0.5283154606819153,
+      "loss_ee": 0.7723551869392395,
+      "loss_delivery": 0.35645291954278946,
+      "loss_biodist": 0.3404483631253242,
+      "loss_toxic": 0.05061229532584548
     },
     {
-      "loss": 2.1527106881141664,
-      "loss_size": 0.19257416054606438,
-      "loss_pdi": 0.5191590428352356,
-      "loss_ee": 0.783897054195404,
-      "loss_delivery": 0.29573799669742584,
-      "loss_biodist": 0.3145490542054176,
-      "loss_toxic": 0.046793402079492806
+      "loss": 2.1428971529006957,
+      "loss_size": 0.19335013553500174,
+      "loss_pdi": 0.5021985083818435,
+      "loss_ee": 0.7642539083957672,
+      "loss_delivery": 0.31821031123399734,
+      "loss_biodist": 0.32588216066360476,
+      "loss_toxic": 0.03900211993604898
     },
     {
-      "loss": 2.0622685074806215,
-      "loss_size": 0.2051038146018982,
-      "loss_pdi": 0.49145313203334806,
-      "loss_ee": 0.7647641122341156,
-      "loss_delivery": 0.2876307189464569,
-      "loss_biodist": 0.27712231278419497,
-      "loss_toxic": 0.03619444826617837
+      "loss": 1.9874909400939942,
+      "loss_size": 0.1736245721578598,
+      "loss_pdi": 0.46206980347633364,
+      "loss_ee": 0.7373365700244904,
+      "loss_delivery": 0.29703493416309357,
+      "loss_biodist": 0.2863417714834213,
+      "loss_toxic": 0.031083252932876348
     },
     {
-      "loss": 1.9519578456878661,
-      "loss_size": 0.17994399815797807,
-      "loss_pdi": 0.4814375311136246,
-      "loss_ee": 0.733842009305954,
-      "loss_delivery": 0.28253656476736067,
-      "loss_biodist": 0.24782671630382538,
-      "loss_toxic": 0.026371066551655532
+      "loss": 1.9297520160675048,
+      "loss_size": 0.1635374441742897,
+      "loss_pdi": 0.4737923800945282,
+      "loss_ee": 0.7171129584312439,
+      "loss_delivery": 0.28808903992176055,
+      "loss_biodist": 0.25874830335378646,
+      "loss_toxic": 0.028471904620528222
     },
     {
-      "loss": 1.935675847530365,
-      "loss_size": 0.1704096481204033,
-      "loss_pdi": 0.47338791787624357,
-      "loss_ee": 0.7182988226413727,
-      "loss_delivery": 0.3093330509960651,
-      "loss_biodist": 0.2340244770050049,
-      "loss_toxic": 0.030221952823922038
+      "loss": 1.8647576332092286,
+      "loss_size": 0.14790172204375268,
+      "loss_pdi": 0.4427785277366638,
+      "loss_ee": 0.7089932143688202,
+      "loss_delivery": 0.30143058970570563,
+      "loss_biodist": 0.24234647750854493,
+      "loss_toxic": 0.021307120053097605
     },
     {
-      "loss": 1.888454306125641,
-      "loss_size": 0.17727438509464263,
-      "loss_pdi": 0.46344051957130433,
-      "loss_ee": 0.7103636503219605,
-      "loss_delivery": 0.30027762055397034,
-      "loss_biodist": 0.2190815806388855,
-      "loss_toxic": 0.018016549991443753
+      "loss": 1.7996623039245605,
+      "loss_size": 0.1429538145661354,
+      "loss_pdi": 0.45114057660102846,
+      "loss_ee": 0.681770408153534,
+      "loss_delivery": 0.2735618159174919,
+      "loss_biodist": 0.2338838443160057,
+      "loss_toxic": 0.01635184111073613
     },
     {
-      "loss": 1.8231052160263062,
-      "loss_size": 0.1548917345702648,
-      "loss_pdi": 0.4576862633228302,
-      "loss_ee": 0.7034903109073639,
-      "loss_delivery": 0.29063438922166823,
-      "loss_biodist": 0.19972888082265855,
-      "loss_toxic": 0.016673638485372066
+      "loss": 1.7303769707679748,
+      "loss_size": 0.13725369721651076,
+      "loss_pdi": 0.43492600619792937,
+      "loss_ee": 0.6648448914289474,
+      "loss_delivery": 0.2714417055249214,
+      "loss_biodist": 0.20898159295320512,
+      "loss_toxic": 0.012929048202931882
     },
     {
-      "loss": 1.756770372390747,
-      "loss_size": 0.15216425359249114,
-      "loss_pdi": 0.429460334777832,
-      "loss_ee": 0.6776757568120957,
-      "loss_delivery": 0.2867794781923294,
-      "loss_biodist": 0.19385820478200913,
-      "loss_toxic": 0.01683232020586729
+      "loss": 1.702065145969391,
+      "loss_size": 0.1783118523657322,
+      "loss_pdi": 0.4118753671646118,
+      "loss_ee": 0.640222480893135,
+      "loss_delivery": 0.2610591858625412,
+      "loss_biodist": 0.20058825612068176,
+      "loss_toxic": 0.01000797227025032
     },
     {
-      "loss": 1.7031532883644105,
-      "loss_size": 0.15495768785476685,
-      "loss_pdi": 0.43094243109226227,
-      "loss_ee": 0.6677232623100281,
-      "loss_delivery": 0.26152765452861787,
-      "loss_biodist": 0.17694738060235976,
-      "loss_toxic": 0.011054831324145198
+      "loss": 1.6243244886398316,
+      "loss_size": 0.1371393844485283,
+      "loss_pdi": 0.3978125751018524,
+      "loss_ee": 0.6315451622009277,
+      "loss_delivery": 0.2618463449180126,
+      "loss_biodist": 0.18574777096509934,
+      "loss_toxic": 0.010233237966895103
     },
     {
-      "loss": 1.679247748851776,
-      "loss_size": 0.15252191424369813,
-      "loss_pdi": 0.40398688316345216,
-      "loss_ee": 0.6563315153121948,
-      "loss_delivery": 0.2827941685914993,
-      "loss_biodist": 0.17421896755695343,
-      "loss_toxic": 0.009394321008585393
+      "loss": 1.645119547843933,
+      "loss_size": 0.13622624576091766,
+      "loss_pdi": 0.4013118803501129,
+      "loss_ee": 0.639850401878357,
+      "loss_delivery": 0.2615354858338833,
+      "loss_biodist": 0.19717498123645782,
+      "loss_toxic": 0.009020529384724797
     },
     {
-      "loss": 1.6231786251068114,
-      "loss_size": 0.14685654938220977,
-      "loss_pdi": 0.3987069964408875,
-      "loss_ee": 0.6459777146577835,
-      "loss_delivery": 0.2517095260322094,
-      "loss_biodist": 0.1695146732032299,
-      "loss_toxic": 0.010413196869194508
+      "loss": 1.5792422771453858,
+      "loss_size": 0.12063037976622581,
+      "loss_pdi": 0.40477685928344725,
+      "loss_ee": 0.6168571084737777,
+      "loss_delivery": 0.23877703920006751,
+      "loss_biodist": 0.1887524366378784,
+      "loss_toxic": 0.009448455832898616
     },
     {
-      "loss": 1.5647669196128846,
-      "loss_size": 0.12480136081576347,
-      "loss_pdi": 0.40768158435821533,
-      "loss_ee": 0.6228045016527176,
-      "loss_delivery": 0.23313914462924004,
-      "loss_biodist": 0.1658004455268383,
-      "loss_toxic": 0.0105398821644485
+      "loss": 1.5701380014419555,
+      "loss_size": 0.12370488420128822,
+      "loss_pdi": 0.3944096490740776,
+      "loss_ee": 0.6204680263996124,
+      "loss_delivery": 0.2499392546713352,
+      "loss_biodist": 0.1741167649626732,
+      "loss_toxic": 0.00749938020016998
     },
     {
-      "loss": 1.543732750415802,
-      "loss_size": 0.12116749435663224,
-      "loss_pdi": 0.3942162901163101,
-      "loss_ee": 0.6175289869308471,
-      "loss_delivery": 0.2506958607584238,
-      "loss_biodist": 0.15138662829995156,
-      "loss_toxic": 0.008737476798705757
+      "loss": 1.5445807576179504,
+      "loss_size": 0.12085893377661705,
+      "loss_pdi": 0.4022176057100296,
+      "loss_ee": 0.6029386401176453,
+      "loss_delivery": 0.2460342638194561,
+      "loss_biodist": 0.16601160615682603,
+      "loss_toxic": 0.006519717467017472
     },
     {
-      "loss": 1.534558892250061,
-      "loss_size": 0.11713396161794662,
-      "loss_pdi": 0.39247085303068163,
-      "loss_ee": 0.608239871263504,
-      "loss_delivery": 0.24818528145551683,
-      "loss_biodist": 0.15680191665887833,
-      "loss_toxic": 0.011727035511285067
+      "loss": 1.4764926195144654,
+      "loss_size": 0.11393929794430732,
+      "loss_pdi": 0.3614879995584488,
+      "loss_ee": 0.5874974340200424,
+      "loss_delivery": 0.2382828861474991,
+      "loss_biodist": 0.168075630068779,
+      "loss_toxic": 0.00720936032012105
     },
     {
-      "loss": 1.508529245853424,
-      "loss_size": 0.12370343580842018,
-      "loss_pdi": 0.37536335587501524,
-      "loss_ee": 0.5983373761177063,
-      "loss_delivery": 0.2428302437067032,
-      "loss_biodist": 0.16056786775588988,
-      "loss_toxic": 0.0077269634697586295
+      "loss": 1.4663256525993347,
+      "loss_size": 0.10480817258358002,
+      "loss_pdi": 0.3699364930391312,
+      "loss_ee": 0.591068571805954,
+      "loss_delivery": 0.23481545299291612,
+      "loss_biodist": 0.1582734301686287,
+      "loss_toxic": 0.007423530006781221
     },
     {
-      "loss": 1.495661163330078,
-      "loss_size": 0.12558167055249214,
-      "loss_pdi": 0.3853023111820221,
-      "loss_ee": 0.5909300297498703,
-      "loss_delivery": 0.22781415805220603,
-      "loss_biodist": 0.15893371179699897,
-      "loss_toxic": 0.007099285908043385
+      "loss": 1.4797919273376465,
+      "loss_size": 0.11906521767377853,
+      "loss_pdi": 0.3831163257360458,
+      "loss_ee": 0.5810098886489868,
+      "loss_delivery": 0.22465722858905793,
+      "loss_biodist": 0.16469249799847602,
+      "loss_toxic": 0.007250743336044252
     },
     {
-      "loss": 1.4543131768703461,
-      "loss_size": 0.11276061460375786,
-      "loss_pdi": 0.3703819438815117,
-      "loss_ee": 0.5739975512027741,
-      "loss_delivery": 0.24190589040517807,
-      "loss_biodist": 0.14822293519973756,
-      "loss_toxic": 0.007044258969835937
-    },
-    {
-      "loss": 1.4692220211029052,
-      "loss_size": 0.12677552737295628,
-      "loss_pdi": 0.37003436386585237,
-      "loss_ee": 0.578819090127945,
-      "loss_delivery": 0.2372341684997082,
-      "loss_biodist": 0.1498504839837551,
-      "loss_toxic": 0.0065083671128377315
+      "loss": 1.4942118644714355,
+      "loss_size": 0.11249525547027588,
+      "loss_pdi": 0.3718418627977371,
+      "loss_ee": 0.5973137259483338,
+      "loss_delivery": 0.23963096588850022,
+      "loss_biodist": 0.16598810032010078,
+      "loss_toxic": 0.006941930414177478
     }
   ],
   "val": [
     {
-      "loss": 16.045034408569336,
-      "loss_size": 7.810531139373779,
-      "loss_pdi": 1.0906392335891724,
-      "loss_ee": 1.0121623277664185,
-      "loss_delivery": 4.599453926086426,
-      "loss_biodist": 1.121813416481018,
-      "loss_toxic": 0.4104346036911011,
-      "acc_pdi": 0.8888888888888888,
+      "loss": 13.683866500854492,
+      "loss_size": 5.657964706420898,
+      "loss_pdi": 1.1590962409973145,
+      "loss_ee": 1.0155898332595825,
+      "loss_delivery": 4.1429033279418945,
+      "loss_biodist": 1.128843069076538,
+      "loss_toxic": 0.579468846321106,
+      "acc_pdi": 0.7407407407407407,
       "acc_ee": 0.6296296296296297,
       "acc_toxic": 1.0
     },
     {
-      "loss": 8.368906021118164,
-      "loss_size": 1.3932932615280151,
-      "loss_pdi": 0.7484031915664673,
-      "loss_ee": 0.9992449283599854,
-      "loss_delivery": 3.9427361488342285,
-      "loss_biodist": 1.0493215322494507,
-      "loss_toxic": 0.2359071671962738,
+      "loss": 7.161172866821289,
+      "loss_size": 0.1799931526184082,
+      "loss_pdi": 0.8303115963935852,
+      "loss_ee": 0.942605197429657,
+      "loss_delivery": 3.986294984817505,
+      "loss_biodist": 1.022797703742981,
+      "loss_toxic": 0.19917015731334686,
       "acc_pdi": 0.8888888888888888,
       "acc_ee": 0.5925925925925926,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.77295446395874,
-      "loss_size": 0.07529407739639282,
-      "loss_pdi": 0.5075266361236572,
-      "loss_ee": 0.9870877265930176,
-      "loss_delivery": 4.065003395080566,
-      "loss_biodist": 1.0452064275741577,
-      "loss_toxic": 0.09283570945262909,
+      "loss": 6.554836273193359,
+      "loss_size": 0.04609166830778122,
+      "loss_pdi": 0.4924769997596741,
+      "loss_ee": 0.965587317943573,
+      "loss_delivery": 3.978637933731079,
+      "loss_biodist": 1.0135102272033691,
+      "loss_toxic": 0.05853228643536568,
       "acc_pdi": 0.8888888888888888,
       "acc_ee": 0.5925925925925926,
       "acc_toxic": 1.0
     },
     {
-      "loss": 7.080923080444336,
-      "loss_size": 0.06525272130966187,
-      "loss_pdi": 0.4343451261520386,
-      "loss_ee": 1.0161277055740356,
-      "loss_delivery": 4.559023380279541,
-      "loss_biodist": 0.9523851871490479,
-      "loss_toxic": 0.05378875508904457,
+      "loss": 6.843129634857178,
+      "loss_size": 0.07650057226419449,
+      "loss_pdi": 0.43551138043403625,
+      "loss_ee": 0.9353340864181519,
+      "loss_delivery": 4.557775974273682,
+      "loss_biodist": 0.7909315228462219,
+      "loss_toxic": 0.047075945883989334,
       "acc_pdi": 0.8888888888888888,
       "acc_ee": 0.5925925925925926,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.916051864624023,
-      "loss_size": 0.07201710343360901,
-      "loss_pdi": 0.4837420880794525,
-      "loss_ee": 1.0075831413269043,
-      "loss_delivery": 4.492982864379883,
-      "loss_biodist": 0.7963473796844482,
-      "loss_toxic": 0.0633791983127594,
+      "loss": 6.711758613586426,
+      "loss_size": 0.04316325858235359,
+      "loss_pdi": 0.41873815655708313,
+      "loss_ee": 1.0096691846847534,
+      "loss_delivery": 4.517927169799805,
+      "loss_biodist": 0.6788683533668518,
+      "loss_toxic": 0.04339226707816124,
       "acc_pdi": 0.8888888888888888,
       "acc_ee": 0.5925925925925926,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.480402946472168,
-      "loss_size": 0.056911345571279526,
-      "loss_pdi": 0.4535364508628845,
-      "loss_ee": 1.00966477394104,
-      "loss_delivery": 4.184553146362305,
-      "loss_biodist": 0.6998977661132812,
-      "loss_toxic": 0.07583901286125183,
+      "loss": 6.905030250549316,
+      "loss_size": 0.045318666845560074,
+      "loss_pdi": 0.38593801856040955,
+      "loss_ee": 1.0019593238830566,
+      "loss_delivery": 4.807835578918457,
+      "loss_biodist": 0.6247215867042542,
+      "loss_toxic": 0.039257097989320755,
+      "acc_pdi": 0.8888888888888888,
+      "acc_ee": 0.5185185185185185,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 6.417820930480957,
+      "loss_size": 0.05034356936812401,
+      "loss_pdi": 0.4149726331233978,
+      "loss_ee": 0.9869357943534851,
+      "loss_delivery": 4.405001640319824,
+      "loss_biodist": 0.533240556716919,
+      "loss_toxic": 0.02732720412313938,
       "acc_pdi": 0.8888888888888888,
       "acc_ee": 0.5925925925925926,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.857071399688721,
-      "loss_size": 0.04192754626274109,
-      "loss_pdi": 0.44999659061431885,
-      "loss_ee": 1.0004281997680664,
-      "loss_delivery": 4.647970676422119,
-      "loss_biodist": 0.6309637427330017,
-      "loss_toxic": 0.0857851505279541,
+      "loss": 6.608631610870361,
+      "loss_size": 0.05222579464316368,
+      "loss_pdi": 0.4375711679458618,
+      "loss_ee": 1.0041171312332153,
+      "loss_delivery": 4.578192234039307,
+      "loss_biodist": 0.5125234723091125,
+      "loss_toxic": 0.02400212176144123,
       "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.5925925925925926,
+      "acc_ee": 0.5185185185185185,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.3087968826293945,
-      "loss_size": 0.09944896399974823,
-      "loss_pdi": 0.40266019105911255,
-      "loss_ee": 0.9705824255943298,
-      "loss_delivery": 4.246408939361572,
-      "loss_biodist": 0.552162230014801,
-      "loss_toxic": 0.03753397986292839,
-      "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.5925925925925926,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 6.9212870597839355,
-      "loss_size": 0.05665857717394829,
-      "loss_pdi": 0.4618251323699951,
-      "loss_ee": 0.9701217412948608,
-      "loss_delivery": 4.835049152374268,
-      "loss_biodist": 0.5396986603736877,
-      "loss_toxic": 0.057933975011110306,
-      "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.5925925925925926,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 6.144314289093018,
-      "loss_size": 0.16736657917499542,
-      "loss_pdi": 0.43004411458969116,
-      "loss_ee": 0.9552963972091675,
-      "loss_delivery": 4.0627570152282715,
-      "loss_biodist": 0.5059604048728943,
-      "loss_toxic": 0.02288975566625595,
-      "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.5925925925925926,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 7.03376579284668,
-      "loss_size": 0.11119232326745987,
-      "loss_pdi": 0.4463132619857788,
-      "loss_ee": 0.9212661385536194,
-      "loss_delivery": 5.010645866394043,
-      "loss_biodist": 0.5135870575904846,
-      "loss_toxic": 0.03076130710542202,
-      "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.5925925925925926,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 6.251324653625488,
-      "loss_size": 0.16327497363090515,
-      "loss_pdi": 0.4076344966888428,
-      "loss_ee": 0.9357188940048218,
-      "loss_delivery": 4.216032981872559,
-      "loss_biodist": 0.5158465504646301,
-      "loss_toxic": 0.012816602364182472,
-      "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.6296296296296297,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 6.909743309020996,
-      "loss_size": 0.10885120183229446,
-      "loss_pdi": 0.40938591957092285,
-      "loss_ee": 0.8893271684646606,
-      "loss_delivery": 4.9792799949646,
-      "loss_biodist": 0.49506261944770813,
-      "loss_toxic": 0.027835864573717117,
-      "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.6296296296296297,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 6.724283695220947,
-      "loss_size": 0.10787779092788696,
-      "loss_pdi": 0.45569828152656555,
-      "loss_ee": 0.979951798915863,
-      "loss_delivery": 4.641026020050049,
-      "loss_biodist": 0.5102843642234802,
-      "loss_toxic": 0.029445137828588486,
+      "loss": 5.7676777839660645,
+      "loss_size": 0.09589201211929321,
+      "loss_pdi": 0.3261733949184418,
+      "loss_ee": 0.9482788443565369,
+      "loss_delivery": 3.856112003326416,
+      "loss_biodist": 0.5298716425895691,
+      "loss_toxic": 0.011350298300385475,
       "acc_pdi": 0.8888888888888888,
       "acc_ee": 0.5555555555555556,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.498813629150391,
-      "loss_size": 0.14617349207401276,
-      "loss_pdi": 0.3836137652397156,
-      "loss_ee": 0.8910271525382996,
-      "loss_delivery": 4.627362251281738,
-      "loss_biodist": 0.43996545672416687,
-      "loss_toxic": 0.010672268457710743,
+      "loss": 6.920990943908691,
+      "loss_size": 0.05388057231903076,
+      "loss_pdi": 0.39705148339271545,
+      "loss_ee": 0.990842878818512,
+      "loss_delivery": 5.025243282318115,
+      "loss_biodist": 0.4346938133239746,
+      "loss_toxic": 0.019278930500149727,
+      "acc_pdi": 0.8888888888888888,
+      "acc_ee": 0.5555555555555556,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 5.798760890960693,
+      "loss_size": 0.09857960045337677,
+      "loss_pdi": 0.33329641819000244,
+      "loss_ee": 0.9614524245262146,
+      "loss_delivery": 4.000489711761475,
+      "loss_biodist": 0.39874210953712463,
+      "loss_toxic": 0.006200834643095732,
+      "acc_pdi": 0.8888888888888888,
+      "acc_ee": 0.5555555555555556,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 6.575327396392822,
+      "loss_size": 0.054699063301086426,
+      "loss_pdi": 0.33702051639556885,
+      "loss_ee": 0.9436452388763428,
+      "loss_delivery": 4.817119121551514,
+      "loss_biodist": 0.41582298278808594,
+      "loss_toxic": 0.007020703982561827,
+      "acc_pdi": 0.8888888888888888,
+      "acc_ee": 0.5555555555555556,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 5.989306449890137,
+      "loss_size": 0.09009546041488647,
+      "loss_pdi": 0.3044246733188629,
+      "loss_ee": 1.0130207538604736,
+      "loss_delivery": 4.140576362609863,
+      "loss_biodist": 0.4378862977027893,
+      "loss_toxic": 0.0033026484306901693,
+      "acc_pdi": 0.8888888888888888,
+      "acc_ee": 0.5555555555555556,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 6.383339881896973,
+      "loss_size": 0.1530081033706665,
+      "loss_pdi": 0.29700207710266113,
+      "loss_ee": 0.9943283796310425,
+      "loss_delivery": 4.564785480499268,
+      "loss_biodist": 0.37085360288619995,
+      "loss_toxic": 0.003362649120390415,
+      "acc_pdi": 0.8888888888888888,
+      "acc_ee": 0.5185185185185185,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 6.233416557312012,
+      "loss_size": 0.1473817676305771,
+      "loss_pdi": 0.2754640281200409,
+      "loss_ee": 0.9803684949874878,
+      "loss_delivery": 4.443488597869873,
+      "loss_biodist": 0.38424524664878845,
+      "loss_toxic": 0.0024682653602212667,
       "acc_pdi": 0.8888888888888888,
       "acc_ee": 0.5925925925925926,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.863435745239258,
-      "loss_size": 0.1379111111164093,
-      "loss_pdi": 0.5173991322517395,
-      "loss_ee": 1.013482689857483,
-      "loss_delivery": 4.7425923347473145,
-      "loss_biodist": 0.4353649318218231,
-      "loss_toxic": 0.01668536849319935,
+      "loss": 6.094257354736328,
+      "loss_size": 0.10127364844083786,
+      "loss_pdi": 0.2960923910140991,
+      "loss_ee": 1.0121080875396729,
+      "loss_delivery": 4.2689008712768555,
+      "loss_biodist": 0.4132467210292816,
+      "loss_toxic": 0.0026356647722423077,
       "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.48148148148148145,
+      "acc_ee": 0.5555555555555556,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.909954071044922,
-      "loss_size": 0.35534632205963135,
-      "loss_pdi": 0.3974299728870392,
-      "loss_ee": 0.8724083304405212,
-      "loss_delivery": 4.841543674468994,
-      "loss_biodist": 0.42888545989990234,
-      "loss_toxic": 0.014340322464704514,
-      "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.6296296296296297,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 6.718777179718018,
-      "loss_size": 0.10662573575973511,
-      "loss_pdi": 0.4000368118286133,
-      "loss_ee": 0.925907552242279,
-      "loss_delivery": 4.855783939361572,
-      "loss_biodist": 0.4174629747867584,
-      "loss_toxic": 0.012960433959960938,
+      "loss": 6.0315470695495605,
+      "loss_size": 0.13236114382743835,
+      "loss_pdi": 0.29554903507232666,
+      "loss_ee": 0.9912998080253601,
+      "loss_delivery": 4.2240777015686035,
+      "loss_biodist": 0.3861069977283478,
+      "loss_toxic": 0.002152827335521579,
       "acc_pdi": 0.8888888888888888,
       "acc_ee": 0.5925925925925926,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.800571918487549,
-      "loss_size": 0.20526628196239471,
-      "loss_pdi": 0.42215225100517273,
-      "loss_ee": 0.9293471574783325,
-      "loss_delivery": 4.7998528480529785,
-      "loss_biodist": 0.42879122495651245,
-      "loss_toxic": 0.015162378549575806,
+      "loss": 6.13291597366333,
+      "loss_size": 0.10603927820920944,
+      "loss_pdi": 0.30880627036094666,
+      "loss_ee": 1.0417256355285645,
+      "loss_delivery": 4.337818622589111,
+      "loss_biodist": 0.33598482608795166,
+      "loss_toxic": 0.002541647758334875,
       "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.5925925925925926,
+      "acc_ee": 0.5555555555555556,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.873208522796631,
-      "loss_size": 0.1934625804424286,
-      "loss_pdi": 0.45854493975639343,
-      "loss_ee": 0.9394232630729675,
-      "loss_delivery": 4.830544471740723,
-      "loss_biodist": 0.43078526854515076,
-      "loss_toxic": 0.020448585972189903,
+      "loss": 5.918347358703613,
+      "loss_size": 0.11423231661319733,
+      "loss_pdi": 0.2779754102230072,
+      "loss_ee": 1.023812174797058,
+      "loss_delivery": 4.137387275695801,
+      "loss_biodist": 0.36283549666404724,
+      "loss_toxic": 0.002104171784594655,
       "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.6296296296296297,
+      "acc_ee": 0.5185185185185185,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.790691375732422,
-      "loss_size": 0.18806225061416626,
-      "loss_pdi": 0.43740272521972656,
-      "loss_ee": 0.9507966637611389,
-      "loss_delivery": 4.783614158630371,
-      "loss_biodist": 0.4198465943336487,
-      "loss_toxic": 0.010969163849949837,
-      "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.5925925925925926,
+      "loss": 6.354115962982178,
+      "loss_size": 0.1212025135755539,
+      "loss_pdi": 0.2848753333091736,
+      "loss_ee": 1.031553030014038,
+      "loss_delivery": 4.554471969604492,
+      "loss_biodist": 0.3598195016384125,
+      "loss_toxic": 0.0021939175203442574,
+      "acc_pdi": 0.8518518518518519,
+      "acc_ee": 0.5555555555555556,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.847814083099365,
-      "loss_size": 0.22717446088790894,
-      "loss_pdi": 0.42288023233413696,
-      "loss_ee": 0.9286114573478699,
-      "loss_delivery": 4.81552791595459,
-      "loss_biodist": 0.43471890687942505,
-      "loss_toxic": 0.01890140399336815,
-      "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.5925925925925926,
+      "loss": 5.881389141082764,
+      "loss_size": 0.1030399352312088,
+      "loss_pdi": 0.2791188657283783,
+      "loss_ee": 1.0205037593841553,
+      "loss_delivery": 4.111578464508057,
+      "loss_biodist": 0.3657107949256897,
+      "loss_toxic": 0.0014369667042046785,
+      "acc_pdi": 0.8518518518518519,
+      "acc_ee": 0.5555555555555556,
       "acc_toxic": 1.0
     },
     {
-      "loss": 7.032270431518555,
-      "loss_size": 0.24087895452976227,
-      "loss_pdi": 0.46614596247673035,
-      "loss_ee": 0.9626513123512268,
-      "loss_delivery": 4.897539138793945,
-      "loss_biodist": 0.44117382168769836,
-      "loss_toxic": 0.023880867287516594,
-      "acc_pdi": 0.6666666666666666,
-      "acc_ee": 0.6296296296296297,
+      "loss": 6.1028852462768555,
+      "loss_size": 0.10241233557462692,
+      "loss_pdi": 0.300007700920105,
+      "loss_ee": 1.0756882429122925,
+      "loss_delivery": 4.258440971374512,
+      "loss_biodist": 0.3646480441093445,
+      "loss_toxic": 0.0016881312476471066,
+      "acc_pdi": 0.8888888888888888,
+      "acc_ee": 0.5185185185185185,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.840210914611816,
-      "loss_size": 0.18546245992183685,
-      "loss_pdi": 0.41406312584877014,
-      "loss_ee": 0.9308509230613708,
-      "loss_delivery": 4.877482891082764,
-      "loss_biodist": 0.41710007190704346,
-      "loss_toxic": 0.015251458622515202,
+      "loss": 6.128824234008789,
+      "loss_size": 0.1437627077102661,
+      "loss_pdi": 0.29325851798057556,
+      "loss_ee": 1.0818182229995728,
+      "loss_delivery": 4.236568450927734,
+      "loss_biodist": 0.3719424605369568,
+      "loss_toxic": 0.001474093529395759,
       "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.6296296296296297,
+      "acc_ee": 0.5185185185185185,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.980130672454834,
-      "loss_size": 0.21011953055858612,
-      "loss_pdi": 0.4370857775211334,
-      "loss_ee": 0.9217703938484192,
-      "loss_delivery": 4.956975936889648,
-      "loss_biodist": 0.4334105849266052,
-      "loss_toxic": 0.020767945796251297,
-      "acc_pdi": 0.8888888888888888,
-      "acc_ee": 0.6296296296296297,
+      "loss": 6.055476188659668,
+      "loss_size": 0.13312266767024994,
+      "loss_pdi": 0.28571468591690063,
+      "loss_ee": 1.066524624824524,
+      "loss_delivery": 4.214193820953369,
+      "loss_biodist": 0.35442692041397095,
+      "loss_toxic": 0.0014939504908397794,
+      "acc_pdi": 0.8518518518518519,
+      "acc_ee": 0.5185185185185185,
       "acc_toxic": 1.0
     }
   ]
diff --git a/models/finetune_cv/fold_0/model.pt b/models/finetune_cv/fold_0/model.pt
index 9e8d77d..943cedc 100644
Binary files a/models/finetune_cv/fold_0/model.pt and b/models/finetune_cv/fold_0/model.pt differ
diff --git a/models/finetune_cv/fold_1/history.json b/models/finetune_cv/fold_1/history.json
index 0873c61..b1306a6 100644
--- a/models/finetune_cv/fold_1/history.json
+++ b/models/finetune_cv/fold_1/history.json
@@ -1,425 +1,425 @@
 {
   "train": [
     {
-      "loss": 22.032494735717773,
-      "loss_size": 16.743953514099122,
-      "loss_pdi": 1.2239572763442994,
-      "loss_ee": 1.0474397182464599,
-      "loss_delivery": 1.2612280070781707,
-      "loss_biodist": 1.1997641563415526,
-      "loss_toxic": 0.5561524748802185
+      "loss": 21.963233947753906,
+      "loss_size": 16.82633171081543,
+      "loss_pdi": 1.2230936765670777,
+      "loss_ee": 1.0703922033309936,
+      "loss_delivery": 1.0690569162368775,
+      "loss_biodist": 1.1534382343292235,
+      "loss_toxic": 0.6209211587905884
     },
     {
-      "loss": 12.498378562927247,
-      "loss_size": 8.05031099319458,
-      "loss_pdi": 1.0179082751274109,
-      "loss_ee": 0.9407736420631408,
-      "loss_delivery": 0.8306634247303009,
-      "loss_biodist": 1.160559320449829,
-      "loss_toxic": 0.4981633126735687
+      "loss": 13.145495796203614,
+      "loss_size": 8.676862716674805,
+      "loss_pdi": 1.0655134558677672,
+      "loss_ee": 0.8999906063079834,
+      "loss_delivery": 0.8303895950317383,
+      "loss_biodist": 1.122723388671875,
+      "loss_toxic": 0.5500160694122315
     },
     {
-      "loss": 6.852249622344971,
-      "loss_size": 2.8376791954040526,
-      "loss_pdi": 0.8515164494514466,
-      "loss_ee": 0.8604451298713685,
-      "loss_delivery": 0.7656278729438781,
-      "loss_biodist": 1.1427257061004639,
-      "loss_toxic": 0.39425510764122007
+      "loss": 7.351448345184326,
+      "loss_size": 3.415665292739868,
+      "loss_pdi": 0.8565655469894409,
+      "loss_ee": 0.7837236523628235,
+      "loss_delivery": 0.8804788589477539,
+      "loss_biodist": 1.011645209789276,
+      "loss_toxic": 0.40336963534355164
     },
     {
-      "loss": 4.40604076385498,
-      "loss_size": 0.6938439428806304,
-      "loss_pdi": 0.6752870678901672,
-      "loss_ee": 0.8080484986305236,
-      "loss_delivery": 0.8313614726066589,
-      "loss_biodist": 1.0828368663787842,
-      "loss_toxic": 0.31466284990310667
+      "loss": 4.39948205947876,
+      "loss_size": 0.9713698267936707,
+      "loss_pdi": 0.6989291191101075,
+      "loss_ee": 0.6805540442466735,
+      "loss_delivery": 0.7624839186668396,
+      "loss_biodist": 0.9798830866813659,
+      "loss_toxic": 0.3062621414661407
     },
     {
-      "loss": 3.4253986835479737,
-      "loss_size": 0.16804716736078262,
-      "loss_pdi": 0.5639804005622864,
-      "loss_ee": 0.6901269078254699,
-      "loss_delivery": 0.6912416338920593,
-      "loss_biodist": 1.0242016911506653,
-      "loss_toxic": 0.28780081272125246
+      "loss": 3.375754451751709,
+      "loss_size": 0.24608666747808455,
+      "loss_pdi": 0.5557448148727417,
+      "loss_ee": 0.6684133768081665,
+      "loss_delivery": 0.7611681580543518,
+      "loss_biodist": 0.919653308391571,
+      "loss_toxic": 0.22468801140785216
     },
     {
-      "loss": 3.2362718105316164,
-      "loss_size": 0.1483217939734459,
-      "loss_pdi": 0.6221840143203735,
-      "loss_ee": 0.6781063556671143,
-      "loss_delivery": 0.6013748198747635,
-      "loss_biodist": 0.9614210963249207,
-      "loss_toxic": 0.22486359924077987
+      "loss": 2.9307605743408205,
+      "loss_size": 0.1106911577284336,
+      "loss_pdi": 0.5004462003707886,
+      "loss_ee": 0.6227471172809601,
+      "loss_delivery": 0.6758030593395233,
+      "loss_biodist": 0.8190896153450012,
+      "loss_toxic": 0.20198351740837098
     },
     {
-      "loss": 3.1696151733398437,
-      "loss_size": 0.1722302332520485,
-      "loss_pdi": 0.48496800661087036,
-      "loss_ee": 0.6616590738296508,
-      "loss_delivery": 0.7679106175899506,
-      "loss_biodist": 0.9077507495880127,
-      "loss_toxic": 0.17509644329547883
+      "loss": 2.731675052642822,
+      "loss_size": 0.13740637749433518,
+      "loss_pdi": 0.4836215674877167,
+      "loss_ee": 0.5896897256374359,
+      "loss_delivery": 0.5866121172904968,
+      "loss_biodist": 0.7556124567985535,
+      "loss_toxic": 0.17873288169503213
     },
     {
-      "loss": 2.6617531299591066,
-      "loss_size": 0.16581893265247344,
-      "loss_pdi": 0.4746619284152985,
-      "loss_ee": 0.630395919084549,
-      "loss_delivery": 0.3917909190058708,
-      "loss_biodist": 0.852727723121643,
-      "loss_toxic": 0.1463577665388584
+      "loss": 2.4887039184570314,
+      "loss_size": 0.12009606957435608,
+      "loss_pdi": 0.4361336886882782,
+      "loss_ee": 0.597134268283844,
+      "loss_delivery": 0.5648026138544082,
+      "loss_biodist": 0.6326960444450378,
+      "loss_toxic": 0.13784122765064238
     },
     {
-      "loss": 2.4909090995788574,
-      "loss_size": 0.11310702562332153,
-      "loss_pdi": 0.43855146765708924,
-      "loss_ee": 0.5929172158241272,
-      "loss_delivery": 0.42325166761875155,
-      "loss_biodist": 0.7738024115562439,
-      "loss_toxic": 0.1492793083190918
+      "loss": 2.1680586099624635,
+      "loss_size": 0.12401954531669616,
+      "loss_pdi": 0.40216060280799865,
+      "loss_ee": 0.5528951227664948,
+      "loss_delivery": 0.42899617552757263,
+      "loss_biodist": 0.5442585527896882,
+      "loss_toxic": 0.1157285787165165
     },
     {
-      "loss": 2.3516653537750245,
-      "loss_size": 0.21800988018512726,
-      "loss_pdi": 0.43570560216903687,
-      "loss_ee": 0.5719826459884644,
-      "loss_delivery": 0.3004884377121925,
-      "loss_biodist": 0.6908805131912231,
-      "loss_toxic": 0.1345983102917671
+      "loss": 2.1059993267059327,
+      "loss_size": 0.13299092650413513,
+      "loss_pdi": 0.38143277168273926,
+      "loss_ee": 0.5274551689624787,
+      "loss_delivery": 0.47739412933588027,
+      "loss_biodist": 0.4953398108482361,
+      "loss_toxic": 0.0913865402340889
     },
     {
-      "loss": 2.1030978202819823,
-      "loss_size": 0.10880238711833953,
-      "loss_pdi": 0.40215051770210264,
-      "loss_ee": 0.5412053823471069,
-      "loss_delivery": 0.2829424023628235,
-      "loss_biodist": 0.651984566450119,
-      "loss_toxic": 0.11601254418492317
+      "loss": 1.9570286750793457,
+      "loss_size": 0.1426382303237915,
+      "loss_pdi": 0.38325140476226804,
+      "loss_ee": 0.49524895548820497,
+      "loss_delivery": 0.42715947031974794,
+      "loss_biodist": 0.4287752747535706,
+      "loss_toxic": 0.07995530962944031
     },
     {
-      "loss": 1.994719886779785,
-      "loss_size": 0.08978431597352028,
-      "loss_pdi": 0.4142456531524658,
-      "loss_ee": 0.532235836982727,
-      "loss_delivery": 0.2747663021087646,
-      "loss_biodist": 0.5863585829734802,
-      "loss_toxic": 0.09732922576367856
+      "loss": 1.8469573497772216,
+      "loss_size": 0.14165955781936646,
+      "loss_pdi": 0.36685559153556824,
+      "loss_ee": 0.4988661766052246,
+      "loss_delivery": 0.36661114990711213,
+      "loss_biodist": 0.39747334718704225,
+      "loss_toxic": 0.07549156174063683
     },
     {
-      "loss": 1.9875550031661988,
-      "loss_size": 0.13663371056318283,
-      "loss_pdi": 0.3811588704586029,
-      "loss_ee": 0.5030429780483245,
-      "loss_delivery": 0.2990179345011711,
-      "loss_biodist": 0.5580029547214508,
-      "loss_toxic": 0.10969849154353142
+      "loss": 1.6980855226516725,
+      "loss_size": 0.11332993358373641,
+      "loss_pdi": 0.350938493013382,
+      "loss_ee": 0.47553136944770813,
+      "loss_delivery": 0.30049399137496946,
+      "loss_biodist": 0.3953311860561371,
+      "loss_toxic": 0.062460555136203764
     },
     {
-      "loss": 1.9142925500869752,
-      "loss_size": 0.12668041437864302,
-      "loss_pdi": 0.4271237254142761,
-      "loss_ee": 0.5021511077880859,
-      "loss_delivery": 0.2249750167131424,
-      "loss_biodist": 0.5295105099678039,
-      "loss_toxic": 0.1038517564535141
+      "loss": 1.743706512451172,
+      "loss_size": 0.12467859983444214,
+      "loss_pdi": 0.3706244468688965,
+      "loss_ee": 0.4802402436733246,
+      "loss_delivery": 0.36484516113996507,
+      "loss_biodist": 0.3557030588388443,
+      "loss_toxic": 0.04761496149003506
     },
     {
-      "loss": 1.81255943775177,
-      "loss_size": 0.1435435712337494,
-      "loss_pdi": 0.3587616056203842,
-      "loss_ee": 0.5099679112434388,
-      "loss_delivery": 0.24120242595672609,
-      "loss_biodist": 0.48327294588088987,
-      "loss_toxic": 0.07581093870103359
+      "loss": 1.7470735549926757,
+      "loss_size": 0.10215002745389938,
+      "loss_pdi": 0.3553147315979004,
+      "loss_ee": 0.4548905730247498,
+      "loss_delivery": 0.4480485826730728,
+      "loss_biodist": 0.3265932142734528,
+      "loss_toxic": 0.06007647253572941
     },
     {
-      "loss": 1.7264988660812377,
-      "loss_size": 0.1290398582816124,
-      "loss_pdi": 0.3462284058332443,
-      "loss_ee": 0.48116588592529297,
-      "loss_delivery": 0.22381204068660737,
-      "loss_biodist": 0.47818992137908933,
-      "loss_toxic": 0.06806278452277184
+      "loss": 1.7687433004379272,
+      "loss_size": 0.10528398901224137,
+      "loss_pdi": 0.35497177839279176,
+      "loss_ee": 0.4946293234825134,
+      "loss_delivery": 0.44853600263595583,
+      "loss_biodist": 0.3113987982273102,
+      "loss_toxic": 0.053923492506146434
     },
     {
-      "loss": 1.6520193338394165,
-      "loss_size": 0.11937985867261887,
-      "loss_pdi": 0.34835702180862427,
-      "loss_ee": 0.4374507278203964,
-      "loss_delivery": 0.24701516777276994,
-      "loss_biodist": 0.43337869048118594,
-      "loss_toxic": 0.06643788442015648
+      "loss": 1.573294997215271,
+      "loss_size": 0.11145550012588501,
+      "loss_pdi": 0.33941014409065245,
+      "loss_ee": 0.42823529839515684,
+      "loss_delivery": 0.34292849004268644,
+      "loss_biodist": 0.3095307767391205,
+      "loss_toxic": 0.04173475466668606
     },
     {
-      "loss": 1.5883747339248657,
-      "loss_size": 0.09780682176351548,
-      "loss_pdi": 0.36214256286621094,
-      "loss_ee": 0.43477209806442263,
-      "loss_delivery": 0.23339744359254838,
-      "loss_biodist": 0.4091412305831909,
-      "loss_toxic": 0.051114612445235255
+      "loss": 1.482050108909607,
+      "loss_size": 0.13211917281150817,
+      "loss_pdi": 0.31831381320953367,
+      "loss_ee": 0.4258797198534012,
+      "loss_delivery": 0.26612227857112886,
+      "loss_biodist": 0.30344046354293824,
+      "loss_toxic": 0.03617466017603874
     },
     {
-      "loss": 1.6899895429611207,
-      "loss_size": 0.10551446527242661,
-      "loss_pdi": 0.3761424541473389,
-      "loss_ee": 0.46849397420883176,
-      "loss_delivery": 0.27737232744693757,
-      "loss_biodist": 0.4076103329658508,
-      "loss_toxic": 0.05485602542757988
+      "loss": 1.5079625368118286,
+      "loss_size": 0.1129397764801979,
+      "loss_pdi": 0.3118207275867462,
+      "loss_ee": 0.4255594819784164,
+      "loss_delivery": 0.30544502288103104,
+      "loss_biodist": 0.31328115463256834,
+      "loss_toxic": 0.03891638442873955
     },
     {
-      "loss": 1.5929443359375,
-      "loss_size": 0.11571613550186158,
-      "loss_pdi": 0.352296257019043,
-      "loss_ee": 0.4432097375392914,
-      "loss_delivery": 0.2575030043721199,
-      "loss_biodist": 0.37044936418533325,
-      "loss_toxic": 0.05376977995038033
+      "loss": 1.4902493238449097,
+      "loss_size": 0.09879767149686813,
+      "loss_pdi": 0.333440762758255,
+      "loss_ee": 0.430321592092514,
+      "loss_delivery": 0.3070627197623253,
+      "loss_biodist": 0.28984564244747163,
+      "loss_toxic": 0.030780918896198273
     }
   ],
   "val": [
     {
-      "loss": 22.81237284342448,
-      "loss_size": 13.806465148925781,
-      "loss_pdi": 1.229181448618571,
-      "loss_ee": 1.027739703655243,
-      "loss_delivery": 5.068911023437977,
-      "loss_biodist": 1.1279457012812297,
-      "loss_toxic": 0.5521289308865865,
-      "acc_pdi": 0.6526315789473685,
-      "acc_ee": 0.6947368421052632,
+      "loss": 24.328961690266926,
+      "loss_size": 15.672358830769857,
+      "loss_pdi": 1.268057902654012,
+      "loss_ee": 1.0569811463356018,
+      "loss_delivery": 4.617272272706032,
+      "loss_biodist": 1.0806464751561482,
+      "loss_toxic": 0.633646289507548,
+      "acc_pdi": 0.6105263157894737,
+      "acc_ee": 0.5894736842105263,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 15.498573303222656,
-      "loss_size": 6.912943522135417,
-      "loss_pdi": 1.151400883992513,
-      "loss_ee": 0.9848727782567342,
-      "loss_delivery": 4.852191311617692,
-      "loss_biodist": 1.106157898902893,
-      "loss_toxic": 0.4910069803396861,
+      "loss": 17.03301429748535,
+      "loss_size": 8.649629751841227,
+      "loss_pdi": 1.165820797284444,
+      "loss_ee": 0.9437925020853678,
+      "loss_delivery": 4.629274984200795,
+      "loss_biodist": 1.0683060089747112,
+      "loss_toxic": 0.5761909882227579,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6736842105263158,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 10.799206574757894,
-      "loss_size": 2.5314422051111856,
-      "loss_pdi": 1.0411948959032695,
-      "loss_ee": 0.9279763499895731,
-      "loss_delivery": 4.807806923985481,
-      "loss_biodist": 1.0742538372675579,
-      "loss_toxic": 0.4165322283903758,
+      "loss": 11.635572751363119,
+      "loss_size": 3.504341204961141,
+      "loss_pdi": 1.0855141083399455,
+      "loss_ee": 0.8674407601356506,
+      "loss_delivery": 4.705501407384872,
+      "loss_biodist": 1.0376905004183452,
+      "loss_toxic": 0.43508487939834595,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6736842105263158,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 8.901862303415934,
-      "loss_size": 0.7716620067755381,
-      "loss_pdi": 1.01361749569575,
-      "loss_ee": 0.8802629311879476,
-      "loss_delivery": 4.857094804445903,
-      "loss_biodist": 1.0347116986910503,
-      "loss_toxic": 0.34451337655385333,
+      "loss": 9.058362166086832,
+      "loss_size": 1.0461570421854656,
+      "loss_pdi": 1.070031762123108,
+      "loss_ee": 0.8463932275772095,
+      "loss_delivery": 4.791346887747447,
+      "loss_biodist": 0.9781110286712646,
+      "loss_toxic": 0.32632239659627277,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6736842105263158,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 8.569346030553183,
-      "loss_size": 0.41188229247927666,
-      "loss_pdi": 1.0502839088439941,
-      "loss_ee": 0.8726372222105662,
-      "loss_delivery": 4.940298028290272,
-      "loss_biodist": 1.0011842250823975,
-      "loss_toxic": 0.29306065539518994,
+      "loss": 8.418675899505615,
+      "loss_size": 0.3764288102587064,
+      "loss_pdi": 1.0916812817255657,
+      "loss_ee": 0.8714254101117452,
+      "loss_delivery": 4.8696667949358625,
+      "loss_biodist": 0.9307892719904581,
+      "loss_toxic": 0.278684730331103,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6736842105263158,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 9.03537893295288,
-      "loss_size": 0.5666823834180832,
-      "loss_pdi": 1.1037296056747437,
-      "loss_ee": 0.9048542082309723,
-      "loss_delivery": 5.218828019996484,
-      "loss_biodist": 0.9743464191754659,
-      "loss_toxic": 0.2669379909833272,
+      "loss": 8.51748021443685,
+      "loss_size": 0.33909208327531815,
+      "loss_pdi": 1.103804111480713,
+      "loss_ee": 0.8707688599824905,
+      "loss_delivery": 5.0624091029167175,
+      "loss_biodist": 0.8743396997451782,
+      "loss_toxic": 0.2670666699608167,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6736842105263158,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 9.051881631215414,
-      "loss_size": 0.5067511014640331,
-      "loss_pdi": 1.0797988673051198,
-      "loss_ee": 0.8918277323246002,
-      "loss_delivery": 5.376374647021294,
-      "loss_biodist": 0.9581284721692404,
-      "loss_toxic": 0.23900071531534195,
+      "loss": 8.701509237289429,
+      "loss_size": 0.38883806640903157,
+      "loss_pdi": 1.0901564558347066,
+      "loss_ee": 0.8219001442193985,
+      "loss_delivery": 5.329233412941297,
+      "loss_biodist": 0.808117667833964,
+      "loss_toxic": 0.2632630293567975,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6736842105263158,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 8.965935707092285,
-      "loss_size": 0.4464708169301351,
-      "loss_pdi": 1.03824187318484,
-      "loss_ee": 0.8565650085608164,
-      "loss_delivery": 5.463352290292581,
-      "loss_biodist": 0.9425446192423502,
-      "loss_toxic": 0.21876097718874613,
+      "loss": 8.602253516515097,
+      "loss_size": 0.399209912866354,
+      "loss_pdi": 1.035650501648585,
+      "loss_ee": 0.8119546920061111,
+      "loss_delivery": 5.297288862367471,
+      "loss_biodist": 0.8003136416276296,
+      "loss_toxic": 0.2578362462421258,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6736842105263158,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 9.351192077000936,
-      "loss_size": 0.578731312106053,
-      "loss_pdi": 1.0086682240168254,
-      "loss_ee": 0.8167769958575567,
-      "loss_delivery": 5.8357385993003845,
-      "loss_biodist": 0.8966569105784098,
-      "loss_toxic": 0.21462025741736093,
+      "loss": 8.610430796941122,
+      "loss_size": 0.3884888291358948,
+      "loss_pdi": 0.9680223266283671,
+      "loss_ee": 0.8063104202349981,
+      "loss_delivery": 5.504999443888664,
+      "loss_biodist": 0.7153328458468119,
+      "loss_toxic": 0.22727691816786924,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6736842105263158,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 9.685971339543661,
-      "loss_size": 0.5851275982956091,
-      "loss_pdi": 0.9895619451999664,
-      "loss_ee": 0.8124870856602987,
-      "loss_delivery": 6.198981747031212,
-      "loss_biodist": 0.8602876861890157,
-      "loss_toxic": 0.2395249493420124,
+      "loss": 8.894750118255615,
+      "loss_size": 0.38015256201227504,
+      "loss_pdi": 0.9849910040696462,
+      "loss_ee": 0.8192636320988337,
+      "loss_delivery": 5.8433875640233355,
+      "loss_biodist": 0.6525928874810537,
+      "loss_toxic": 0.21436312049627304,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6736842105263158,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 9.098868290583292,
-      "loss_size": 0.4912531226873398,
-      "loss_pdi": 0.9610133767127991,
-      "loss_ee": 0.8015020589033762,
-      "loss_delivery": 5.819371705253919,
-      "loss_biodist": 0.8127925594647726,
-      "loss_toxic": 0.21293580221633115,
+      "loss": 8.684672435124716,
+      "loss_size": 0.39142270882924396,
+      "loss_pdi": 0.9926454623540243,
+      "loss_ee": 0.8487897912661234,
+      "loss_delivery": 5.675399616360664,
+      "loss_biodist": 0.5763055086135864,
+      "loss_toxic": 0.2001086367915074,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6736842105263158,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 9.217247327168783,
-      "loss_size": 0.5547616928815842,
-      "loss_pdi": 0.9466870129108429,
-      "loss_ee": 0.80243648091952,
-      "loss_delivery": 5.907406737407048,
-      "loss_biodist": 0.7983624935150146,
-      "loss_toxic": 0.20759239544471106,
+      "loss": 8.468807538350424,
+      "loss_size": 0.37035099665323895,
+      "loss_pdi": 0.9933059811592102,
+      "loss_ee": 0.8365495651960373,
+      "loss_delivery": 5.39086152613163,
+      "loss_biodist": 0.6555034021536509,
+      "loss_toxic": 0.22223659542699656,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6736842105263158,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 9.307894190152487,
-      "loss_size": 0.6155262216925621,
-      "loss_pdi": 0.9505135516325632,
-      "loss_ee": 0.8020251393318176,
-      "loss_delivery": 5.944891105095546,
-      "loss_biodist": 0.7798070808251699,
-      "loss_toxic": 0.21513095125555992,
+      "loss": 8.48760732014974,
+      "loss_size": 0.3547621878484885,
+      "loss_pdi": 1.008083571990331,
+      "loss_ee": 0.8507340376575788,
+      "loss_delivery": 5.329072058200836,
+      "loss_biodist": 0.7051869928836823,
+      "loss_toxic": 0.23976873668531576,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6736842105263158,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 9.434777816136679,
-      "loss_size": 0.5672734752297401,
-      "loss_pdi": 0.9817352195580801,
-      "loss_ee": 0.822968602180481,
-      "loss_delivery": 6.080470234155655,
-      "loss_biodist": 0.745076318581899,
-      "loss_toxic": 0.23725438863039017,
-      "acc_pdi": 0.6105263157894737,
-      "acc_ee": 0.6736842105263158,
-      "acc_toxic": 0.8939393939393939
-    },
-    {
-      "loss": 9.355290253957113,
-      "loss_size": 0.5713960801561674,
-      "loss_pdi": 0.9858902891476949,
-      "loss_ee": 0.8192337850729624,
-      "loss_delivery": 6.011797075470288,
-      "loss_biodist": 0.724964420000712,
-      "loss_toxic": 0.24200843647122383,
-      "acc_pdi": 0.6105263157894737,
-      "acc_ee": 0.6736842105263158,
-      "acc_toxic": 0.8939393939393939
-    },
-    {
-      "loss": 9.306842883427938,
-      "loss_size": 0.6075545425216357,
-      "loss_pdi": 0.9776454170544943,
-      "loss_ee": 0.7862655818462372,
-      "loss_delivery": 5.985511064529419,
-      "loss_biodist": 0.7076093653837839,
-      "loss_toxic": 0.24225737899541855,
-      "acc_pdi": 0.6105263157894737,
-      "acc_ee": 0.6736842105263158,
-      "acc_toxic": 0.8939393939393939
-    },
-    {
-      "loss": 9.252086718877157,
-      "loss_size": 0.6051755348841349,
-      "loss_pdi": 0.9714606404304504,
-      "loss_ee": 0.7621750583251318,
-      "loss_delivery": 5.975689520438512,
-      "loss_biodist": 0.6939358512560526,
-      "loss_toxic": 0.24364992106954256,
-      "acc_pdi": 0.6105263157894737,
-      "acc_ee": 0.6736842105263158,
-      "acc_toxic": 0.8939393939393939
-    },
-    {
-      "loss": 9.275835235913595,
-      "loss_size": 0.5751028036077818,
-      "loss_pdi": 0.9899951120217642,
-      "loss_ee": 0.768888125816981,
-      "loss_delivery": 5.988193516929944,
-      "loss_biodist": 0.6993262469768524,
-      "loss_toxic": 0.25432955101132393,
+      "loss": 8.534255782763163,
+      "loss_size": 0.35214799270033836,
+      "loss_pdi": 1.0083338419596355,
+      "loss_ee": 0.8703259030977885,
+      "loss_delivery": 5.4809657235940294,
+      "loss_biodist": 0.6066243648529053,
+      "loss_toxic": 0.21585797673712173,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6842105263157895,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 9.215376615524292,
-      "loss_size": 0.5498512660463651,
-      "loss_pdi": 0.9988046189149221,
-      "loss_ee": 0.7566283419728279,
-      "loss_delivery": 5.962520445386569,
-      "loss_biodist": 0.6932495137055715,
-      "loss_toxic": 0.2543224884817998,
+      "loss": 8.59092911084493,
+      "loss_size": 0.3520332872867584,
+      "loss_pdi": 0.9944024880727133,
+      "loss_ee": 0.8839219162861506,
+      "loss_delivery": 5.593439628680547,
+      "loss_biodist": 0.562449519832929,
+      "loss_toxic": 0.20468231476843357,
+      "acc_pdi": 0.6105263157894737,
+      "acc_ee": 0.6736842105263158,
+      "acc_toxic": 0.8939393939393939
+    },
+    {
+      "loss": 8.581690510114035,
+      "loss_size": 0.3468632685641448,
+      "loss_pdi": 1.0153752664724986,
+      "loss_ee": 0.884696863591671,
+      "loss_delivery": 5.548932209610939,
+      "loss_biodist": 0.576594889163971,
+      "loss_toxic": 0.20922777770708004,
+      "acc_pdi": 0.6105263157894737,
+      "acc_ee": 0.6736842105263158,
+      "acc_toxic": 0.8939393939393939
+    },
+    {
+      "loss": 8.60028068224589,
+      "loss_size": 0.34553587809205055,
+      "loss_pdi": 1.0314316948254902,
+      "loss_ee": 0.8696443388859431,
+      "loss_delivery": 5.513105024894078,
+      "loss_biodist": 0.6187789390484492,
+      "loss_toxic": 0.2217849005634586,
       "acc_pdi": 0.6105263157894737,
       "acc_ee": 0.6947368421052632,
       "acc_toxic": 0.8939393939393939
     },
     {
-      "loss": 9.17027481396993,
-      "loss_size": 0.5183901910980543,
-      "loss_pdi": 1.0107523500919342,
-      "loss_ee": 0.7627487083276113,
-      "loss_delivery": 5.940715471903483,
-      "loss_biodist": 0.6853441596031189,
-      "loss_toxic": 0.25232408692439395,
+      "loss": 8.721842130025228,
+      "loss_size": 0.3432792164385319,
+      "loss_pdi": 1.044082870086034,
+      "loss_ee": 0.8888355021675428,
+      "loss_delivery": 5.590167284011841,
+      "loss_biodist": 0.6300752113262812,
+      "loss_toxic": 0.22540184513976178,
       "acc_pdi": 0.6105263157894737,
-      "acc_ee": 0.6947368421052632,
+      "acc_ee": 0.6736842105263158,
+      "acc_toxic": 0.8939393939393939
+    },
+    {
+      "loss": 8.821967244148254,
+      "loss_size": 0.3423520748813947,
+      "loss_pdi": 1.0627215206623077,
+      "loss_ee": 0.9012102037668228,
+      "loss_delivery": 5.6443866689999895,
+      "loss_biodist": 0.6428664823373159,
+      "loss_toxic": 0.2284308553983768,
+      "acc_pdi": 0.6105263157894737,
+      "acc_ee": 0.6631578947368421,
+      "acc_toxic": 0.8939393939393939
+    },
+    {
+      "loss": 8.798149506251017,
+      "loss_size": 0.3439513569076856,
+      "loss_pdi": 1.074403668443362,
+      "loss_ee": 0.9037297517061234,
+      "loss_delivery": 5.62445667386055,
+      "loss_biodist": 0.6279164751370748,
+      "loss_toxic": 0.22369086369872093,
+      "acc_pdi": 0.6105263157894737,
+      "acc_ee": 0.6631578947368421,
       "acc_toxic": 0.8939393939393939
     }
   ]
diff --git a/models/finetune_cv/fold_1/model.pt b/models/finetune_cv/fold_1/model.pt
index 9e3e48d..e8cd0be 100644
Binary files a/models/finetune_cv/fold_1/model.pt and b/models/finetune_cv/fold_1/model.pt differ
diff --git a/models/finetune_cv/fold_2/history.json b/models/finetune_cv/fold_2/history.json
index a3ad4a4..5ca8634 100644
--- a/models/finetune_cv/fold_2/history.json
+++ b/models/finetune_cv/fold_2/history.json
@@ -1,467 +1,530 @@
 {
   "train": [
     {
-      "loss": 23.856885592142742,
-      "loss_size": 17.954859733581543,
-      "loss_pdi": 1.3681265513102214,
-      "loss_ee": 1.068708558877309,
-      "loss_delivery": 1.6815399676561356,
-      "loss_biodist": 1.1385973294576008,
-      "loss_toxic": 0.6450536052385966
+      "loss": 22.972569465637207,
+      "loss_size": 17.22947899500529,
+      "loss_pdi": 1.3510672052701314,
+      "loss_ee": 1.0506827433904011,
+      "loss_delivery": 1.5721548050642014,
+      "loss_biodist": 1.1304322481155396,
+      "loss_toxic": 0.6387530366579691
     },
     {
-      "loss": 15.18941100438436,
-      "loss_size": 10.012138843536377,
-      "loss_pdi": 1.2400256196657817,
-      "loss_ee": 0.9122016330560049,
-      "loss_delivery": 1.4215861509243648,
-      "loss_biodist": 1.0704677999019623,
-      "loss_toxic": 0.5329908380905787
+      "loss": 12.718681335449219,
+      "loss_size": 7.335077285766602,
+      "loss_pdi": 1.198062241077423,
+      "loss_ee": 0.97108127673467,
+      "loss_delivery": 1.639556477467219,
+      "loss_biodist": 1.0746847093105316,
+      "loss_toxic": 0.50021959344546
     },
     {
-      "loss": 8.939830541610718,
-      "loss_size": 4.306200265884399,
-      "loss_pdi": 1.0988461375236511,
-      "loss_ee": 0.7973864078521729,
-      "loss_delivery": 1.312243824203809,
-      "loss_biodist": 0.9968815743923187,
-      "loss_toxic": 0.4282720486323039
+      "loss": 6.867454210917155,
+      "loss_size": 2.3153140544891357,
+      "loss_pdi": 1.0159071584542592,
+      "loss_ee": 0.853508859872818,
+      "loss_delivery": 1.2862873176733653,
+      "loss_biodist": 1.0416639745235443,
+      "loss_toxic": 0.35477257271607715
     },
     {
-      "loss": 5.871116876602173,
-      "loss_size": 1.4981385171413422,
-      "loss_pdi": 0.9578391710917155,
-      "loss_ee": 0.7877674202124277,
-      "loss_delivery": 1.332635521888733,
-      "loss_biodist": 0.9970230559508005,
-      "loss_toxic": 0.29771314313014346
+      "loss": 4.856432318687439,
+      "loss_size": 0.5409951706727346,
+      "loss_pdi": 0.8652523259321848,
+      "loss_ee": 0.7771940131982168,
+      "loss_delivery": 1.413562481602033,
+      "loss_biodist": 0.9977987806002299,
+      "loss_toxic": 0.26162934054931003
     },
     {
-      "loss": 4.818921804428101,
-      "loss_size": 0.6080186615387598,
-      "loss_pdi": 0.8167077898979187,
-      "loss_ee": 0.7845198512077332,
-      "loss_delivery": 1.421961595614751,
-      "loss_biodist": 0.9678046603997549,
-      "loss_toxic": 0.21990922341744104
+      "loss": 4.253215591112773,
+      "loss_size": 0.2641367167234421,
+      "loss_pdi": 0.739859402179718,
+      "loss_ee": 0.7256686190764109,
+      "loss_delivery": 1.4241955528656642,
+      "loss_biodist": 0.8935903211434683,
+      "loss_toxic": 0.2057649294535319
     },
     {
-      "loss": 4.394984285036723,
-      "loss_size": 0.3149821311235428,
-      "loss_pdi": 0.7380032440026602,
-      "loss_ee": 0.7671190500259399,
-      "loss_delivery": 1.447837049762408,
-      "loss_biodist": 0.9290279944737753,
-      "loss_toxic": 0.19801472003261247
+      "loss": 3.8961705764134726,
+      "loss_size": 0.2962125514944394,
+      "loss_pdi": 0.682400623957316,
+      "loss_ee": 0.6820215880870819,
+      "loss_delivery": 1.2787245536843936,
+      "loss_biodist": 0.7812575101852417,
+      "loss_toxic": 0.1755537080268065
     },
     {
-      "loss": 3.884276866912842,
-      "loss_size": 0.26142628739277524,
-      "loss_pdi": 0.6927057504653931,
-      "loss_ee": 0.7489989002545675,
-      "loss_delivery": 1.1156622817118962,
-      "loss_biodist": 0.8776024182637533,
-      "loss_toxic": 0.1878812573850155
+      "loss": 3.4790991942087808,
+      "loss_size": 0.3047281603018443,
+      "loss_pdi": 0.6409291823705038,
+      "loss_ee": 0.6178905169169108,
+      "loss_delivery": 1.1121559316913288,
+      "loss_biodist": 0.6434484819571177,
+      "loss_toxic": 0.15994682783881822
     },
     {
-      "loss": 3.709937810897827,
-      "loss_size": 0.37482741723457974,
-      "loss_pdi": 0.6356837352116903,
-      "loss_ee": 0.7201081812381744,
-      "loss_delivery": 1.0413622185587883,
-      "loss_biodist": 0.769838293393453,
-      "loss_toxic": 0.16811783549686274
+      "loss": 3.2075613339742026,
+      "loss_size": 0.3421506683031718,
+      "loss_pdi": 0.5879766543706259,
+      "loss_ee": 0.5811398377021154,
+      "loss_delivery": 1.0462109719713528,
+      "loss_biodist": 0.520307645201683,
+      "loss_toxic": 0.12977550799647966
     },
     {
-      "loss": 3.3223368724187217,
-      "loss_size": 0.2915526789923509,
-      "loss_pdi": 0.5975265900293986,
-      "loss_ee": 0.657735288143158,
-      "loss_delivery": 0.915939765671889,
-      "loss_biodist": 0.6836796899636587,
-      "loss_toxic": 0.17590284595886865
+      "loss": 2.861353278160095,
+      "loss_size": 0.2742840300003688,
+      "loss_pdi": 0.5437282969554266,
+      "loss_ee": 0.5531725088755289,
+      "loss_delivery": 0.9213679246604443,
+      "loss_biodist": 0.4499489863713582,
+      "loss_toxic": 0.11885150956610839
     },
     {
-      "loss": 3.22677751382192,
-      "loss_size": 0.30531836052735645,
-      "loss_pdi": 0.5498206416765848,
-      "loss_ee": 0.6144644021987915,
-      "loss_delivery": 1.0515401139855385,
-      "loss_biodist": 0.5715167572100958,
-      "loss_toxic": 0.13411726988852024
+      "loss": 2.6909215847651162,
+      "loss_size": 0.23881135260065398,
+      "loss_pdi": 0.5229279547929764,
+      "loss_ee": 0.5285524874925613,
+      "loss_delivery": 0.8911051253477732,
+      "loss_biodist": 0.4015616128842036,
+      "loss_toxic": 0.10796305599311988
     },
     {
-      "loss": 2.898585855960846,
-      "loss_size": 0.26108303914467496,
-      "loss_pdi": 0.5470793843269348,
-      "loss_ee": 0.5787178675333658,
-      "loss_delivery": 0.9020876735448837,
-      "loss_biodist": 0.47204887370268506,
-      "loss_toxic": 0.13756892209251723
+      "loss": 2.5927247206370034,
+      "loss_size": 0.27356760079662007,
+      "loss_pdi": 0.5166990955670675,
+      "loss_ee": 0.5059170673290888,
+      "loss_delivery": 0.8377179056406021,
+      "loss_biodist": 0.3519642899433772,
+      "loss_toxic": 0.10685871541500092
     },
     {
-      "loss": 2.6754438877105713,
-      "loss_size": 0.27526700248320896,
-      "loss_pdi": 0.5229905943075815,
-      "loss_ee": 0.5479903370141983,
-      "loss_delivery": 0.7761634774506092,
-      "loss_biodist": 0.44562476873397827,
-      "loss_toxic": 0.10740765929222107
+      "loss": 2.3971973856290183,
+      "loss_size": 0.2688147674004237,
+      "loss_pdi": 0.4851151605447133,
+      "loss_ee": 0.47870688637097675,
+      "loss_delivery": 0.7584750155607859,
+      "loss_biodist": 0.3166690344611804,
+      "loss_toxic": 0.08941652067005634
     },
     {
-      "loss": 2.4977574348449707,
-      "loss_size": 0.2336499529580275,
-      "loss_pdi": 0.4902036637067795,
-      "loss_ee": 0.5169308086236318,
-      "loss_delivery": 0.7469637009004751,
-      "loss_biodist": 0.40839699904123944,
-      "loss_toxic": 0.10161229533453782
+      "loss": 2.2271180947621665,
+      "loss_size": 0.2559296215573947,
+      "loss_pdi": 0.467803418636322,
+      "loss_ee": 0.4819647620121638,
+      "loss_delivery": 0.6487737223505974,
+      "loss_biodist": 0.2930952211221059,
+      "loss_toxic": 0.079551310899357
     },
     {
-      "loss": 2.384280482927958,
-      "loss_size": 0.2612900485595067,
-      "loss_pdi": 0.48895320296287537,
-      "loss_ee": 0.49676452577114105,
-      "loss_delivery": 0.6831860815485319,
-      "loss_biodist": 0.36347728967666626,
-      "loss_toxic": 0.090609318887194
+      "loss": 2.1467134952545166,
+      "loss_size": 0.2658323546250661,
+      "loss_pdi": 0.47287177046140033,
+      "loss_ee": 0.4580538024504979,
+      "loss_delivery": 0.6110207016269366,
+      "loss_biodist": 0.26590356479088467,
+      "loss_toxic": 0.07303123424450557
     },
     {
-      "loss": 2.3147188226381936,
-      "loss_size": 0.25488365814089775,
-      "loss_pdi": 0.4683869779109955,
-      "loss_ee": 0.4875288059314092,
-      "loss_delivery": 0.6712647005915642,
-      "loss_biodist": 0.35311167935530346,
-      "loss_toxic": 0.0795429985349377
+      "loss": 2.0699684421221414,
+      "loss_size": 0.23655260602633157,
+      "loss_pdi": 0.46446068088213605,
+      "loss_ee": 0.43884341915448505,
+      "loss_delivery": 0.5945644030968348,
+      "loss_biodist": 0.26856863250335056,
+      "loss_toxic": 0.06697871504972379
     },
     {
-      "loss": 2.2636735240618386,
-      "loss_size": 0.25234917054573697,
-      "loss_pdi": 0.48236118257045746,
-      "loss_ee": 0.4705241521199544,
-      "loss_delivery": 0.6308227330446243,
-      "loss_biodist": 0.3382392128308614,
-      "loss_toxic": 0.0893770344555378
+      "loss": 2.012367367744446,
+      "loss_size": 0.20358355715870857,
+      "loss_pdi": 0.44864421089490253,
+      "loss_ee": 0.4260970900456111,
+      "loss_delivery": 0.6111055202782154,
+      "loss_biodist": 0.24829111248254776,
+      "loss_toxic": 0.07464585608492295
     },
     {
-      "loss": 2.116434315840403,
-      "loss_size": 0.24041289339462915,
-      "loss_pdi": 0.45973017315069836,
-      "loss_ee": 0.4558122158050537,
-      "loss_delivery": 0.5561455090840658,
-      "loss_biodist": 0.3245606869459152,
-      "loss_toxic": 0.07977284273753564
+      "loss": 1.9354575673739116,
+      "loss_size": 0.19155597686767578,
+      "loss_pdi": 0.43001438677310944,
+      "loss_ee": 0.4029633104801178,
+      "loss_delivery": 0.5866967861851057,
+      "loss_biodist": 0.26284457246462506,
+      "loss_toxic": 0.06138256782044967
     },
     {
-      "loss": 2.116472323735555,
-      "loss_size": 0.22036718018352985,
-      "loss_pdi": 0.4676543176174164,
-      "loss_ee": 0.43863776326179504,
-      "loss_delivery": 0.6074383656183878,
-      "loss_biodist": 0.3068434993426005,
-      "loss_toxic": 0.07553133244315784
+      "loss": 1.9248821139335632,
+      "loss_size": 0.19836385796467462,
+      "loss_pdi": 0.43165912727514905,
+      "loss_ee": 0.4223821411530177,
+      "loss_delivery": 0.5774712382505337,
+      "loss_biodist": 0.23008103668689728,
+      "loss_toxic": 0.06492467441906531
     },
     {
-      "loss": 2.0553239782651267,
-      "loss_size": 0.2092201883594195,
-      "loss_pdi": 0.45897047221660614,
-      "loss_ee": 0.44677118460337323,
-      "loss_delivery": 0.559788204729557,
-      "loss_biodist": 0.31095271309216815,
-      "loss_toxic": 0.06962121867885192
+      "loss": 1.7986130317052205,
+      "loss_size": 0.1977602814634641,
+      "loss_pdi": 0.4213625093301137,
+      "loss_ee": 0.3969506522019704,
+      "loss_delivery": 0.4972396679222584,
+      "loss_biodist": 0.22815552850564322,
+      "loss_toxic": 0.05714430411656698
     },
     {
-      "loss": 1.961161474386851,
-      "loss_size": 0.20881337051590285,
-      "loss_pdi": 0.45118602613608044,
-      "loss_ee": 0.42851509153842926,
-      "loss_delivery": 0.49867390592892963,
-      "loss_biodist": 0.3025246188044548,
-      "loss_toxic": 0.07144851268579562
+      "loss": 1.8008437156677246,
+      "loss_size": 0.20143492271502814,
+      "loss_pdi": 0.4257240394751231,
+      "loss_ee": 0.3939937750498454,
+      "loss_delivery": 0.4996156108876069,
+      "loss_biodist": 0.22945881386597952,
+      "loss_toxic": 0.05061656702309847
     },
     {
-      "loss": 1.959239919980367,
-      "loss_size": 0.2156102918088436,
-      "loss_pdi": 0.4368931899468104,
-      "loss_ee": 0.42987839380900067,
-      "loss_delivery": 0.5220988343159357,
-      "loss_biodist": 0.2822929248213768,
-      "loss_toxic": 0.07246625237166882
+      "loss": 1.8123606244723003,
+      "loss_size": 0.23175274084011713,
+      "loss_pdi": 0.41065867245197296,
+      "loss_ee": 0.38645289838314056,
+      "loss_delivery": 0.5105274474869171,
+      "loss_biodist": 0.22759289046128592,
+      "loss_toxic": 0.04537593169758717
     },
     {
-      "loss": 1.9401288827260335,
-      "loss_size": 0.21518264586726824,
-      "loss_pdi": 0.45129939913749695,
-      "loss_ee": 0.4190533608198166,
-      "loss_delivery": 0.5097754697004954,
-      "loss_biodist": 0.27901960412661236,
-      "loss_toxic": 0.06579839282979567
+      "loss": 1.85766206185023,
+      "loss_size": 0.2237167110045751,
+      "loss_pdi": 0.4198872745037079,
+      "loss_ee": 0.39036936064561206,
+      "loss_delivery": 0.5356200908621153,
+      "loss_biodist": 0.2327907457947731,
+      "loss_toxic": 0.055277835965777435
+    },
+    {
+      "loss": 1.7299150824546814,
+      "loss_size": 0.16866947089632353,
+      "loss_pdi": 0.40182044357061386,
+      "loss_ee": 0.37123599648475647,
+      "loss_delivery": 0.5183743331581354,
+      "loss_biodist": 0.22459317495425543,
+      "loss_toxic": 0.04522162117063999
+    },
+    {
+      "loss": 1.8115381598472595,
+      "loss_size": 0.21021889025966325,
+      "loss_pdi": 0.3938516428073247,
+      "loss_ee": 0.3856282929579417,
+      "loss_delivery": 0.5463737193495035,
+      "loss_biodist": 0.22467835744222006,
+      "loss_toxic": 0.05078731415172418
+    },
+    {
+      "loss": 1.7609570423762004,
+      "loss_size": 0.20672637100021043,
+      "loss_pdi": 0.3868243644634883,
+      "loss_ee": 0.3775654385487239,
+      "loss_delivery": 0.5160359914104143,
+      "loss_biodist": 0.22730938345193863,
+      "loss_toxic": 0.04649555139864484
     }
   ],
   "val": [
     {
-      "loss": 17.15616934640067,
-      "loss_size": 12.466235705784388,
-      "loss_pdi": 1.2282596485955375,
-      "loss_ee": 1.1713778802326746,
-      "loss_delivery": 0.45601305684873034,
-      "loss_biodist": 1.3372918111937386,
-      "loss_toxic": 0.49699102129255024,
-      "acc_pdi": 0.6410256410256411,
-      "acc_ee": 0.4205128205128205,
+      "loss": 19.896042142595565,
+      "loss_size": 14.947636876787458,
+      "loss_pdi": 1.3514722074781145,
+      "loss_ee": 1.0372784308024816,
+      "loss_delivery": 0.5157596128327506,
+      "loss_biodist": 1.3665738276072912,
+      "loss_toxic": 0.6773212381771633,
+      "acc_pdi": 0.22564102564102564,
+      "acc_ee": 0.4512820512820513,
+      "acc_toxic": 0.7073170731707317
+    },
+    {
+      "loss": 10.277108192443848,
+      "loss_size": 5.728530270712716,
+      "loss_pdi": 1.2047701733452933,
+      "loss_ee": 1.013599353177207,
+      "loss_delivery": 0.5329383058207375,
+      "loss_biodist": 1.3288453817367554,
+      "loss_toxic": 0.4684244394302368,
+      "acc_pdi": 0.358974358974359,
+      "acc_ee": 0.4461538461538462,
       "acc_toxic": 1.0
     },
     {
-      "loss": 9.976628576006208,
-      "loss_size": 5.5366509301321845,
-      "loss_pdi": 1.1110572814941406,
-      "loss_ee": 1.2255418130329676,
-      "loss_delivery": 0.4425822538988931,
-      "loss_biodist": 1.2637801681246077,
-      "loss_toxic": 0.3970160186290741,
+      "loss": 4.947442190987723,
+      "loss_size": 0.9055690084184919,
+      "loss_pdi": 0.9325166344642639,
+      "loss_ee": 1.071527932371412,
+      "loss_delivery": 0.5525430504764829,
+      "loss_biodist": 1.2514750446592058,
+      "loss_toxic": 0.23381062703473227,
       "acc_pdi": 0.7076923076923077,
       "acc_ee": 0.4205128205128205,
       "acc_toxic": 1.0
     },
     {
-      "loss": 6.142949717385428,
-      "loss_size": 2.012831313269479,
-      "loss_pdi": 0.9936460341726031,
-      "loss_ee": 1.2630535023553031,
-      "loss_delivery": 0.45085248563970837,
-      "loss_biodist": 1.1773428661482674,
-      "loss_toxic": 0.24522356688976288,
+      "loss": 4.072668756757464,
+      "loss_size": 0.18621658267719404,
+      "loss_pdi": 0.7640595691544669,
+      "loss_ee": 1.2155327456338065,
+      "loss_delivery": 0.5523517067943301,
+      "loss_biodist": 1.214196733066014,
+      "loss_toxic": 0.14031140506267548,
       "acc_pdi": 0.7076923076923077,
       "acc_ee": 0.4205128205128205,
       "acc_toxic": 1.0
     },
     {
-      "loss": 4.523837975093296,
-      "loss_size": 0.6790421732834407,
-      "loss_pdi": 0.8449332543781826,
-      "loss_ee": 1.2828129359654017,
-      "loss_delivery": 0.47022694775036405,
-      "loss_biodist": 1.0964353680610657,
-      "loss_toxic": 0.15038717218807765,
+      "loss": 4.053883859089443,
+      "loss_size": 0.24764748875583922,
+      "loss_pdi": 0.6859285916600909,
+      "loss_ee": 1.3019903557641166,
+      "loss_delivery": 0.5384655041354043,
+      "loss_biodist": 1.191110406603132,
+      "loss_toxic": 0.08874151536396571,
       "acc_pdi": 0.7076923076923077,
       "acc_ee": 0.4205128205128205,
       "acc_toxic": 1.0
     },
     {
-      "loss": 3.914861168180193,
-      "loss_size": 0.23712052990283286,
-      "loss_pdi": 0.739460038287299,
-      "loss_ee": 1.307591336114066,
-      "loss_delivery": 0.48760053728307995,
-      "loss_biodist": 1.0441895723342896,
-      "loss_toxic": 0.09889917501381465,
+      "loss": 3.9850601128169467,
+      "loss_size": 0.18243093735405377,
+      "loss_pdi": 0.6606386282614299,
+      "loss_ee": 1.2955879313605172,
+      "loss_delivery": 0.6222422846726009,
+      "loss_biodist": 1.1586603011403764,
+      "loss_toxic": 0.06550001353025436,
       "acc_pdi": 0.7076923076923077,
       "acc_ee": 0.4205128205128205,
       "acc_toxic": 1.0
     },
     {
-      "loss": 3.827185903276716,
-      "loss_size": 0.19252057054213115,
-      "loss_pdi": 0.6754980896200452,
-      "loss_ee": 1.3178976603916712,
-      "loss_delivery": 0.52304607629776,
-      "loss_biodist": 1.0516272272382463,
-      "loss_toxic": 0.06659629621676036,
+      "loss": 3.6691461631229947,
+      "loss_size": 0.18682856378810747,
+      "loss_pdi": 0.6505168399640492,
+      "loss_ee": 1.2279302733285087,
+      "loss_delivery": 0.5776888344969068,
+      "loss_biodist": 0.963392470564161,
+      "loss_toxic": 0.06278917672378677,
       "acc_pdi": 0.7076923076923077,
       "acc_ee": 0.4205128205128205,
       "acc_toxic": 1.0
     },
     {
-      "loss": 3.7409281730651855,
-      "loss_size": 0.19949970713683537,
-      "loss_pdi": 0.63478513274874,
-      "loss_ee": 1.3695382390703474,
-      "loss_delivery": 0.5511368151221957,
-      "loss_biodist": 0.9442958320890155,
-      "loss_toxic": 0.04167233567152705,
+      "loss": 3.6420267650059293,
+      "loss_size": 0.1952320017984935,
+      "loss_pdi": 0.6510439600263324,
+      "loss_ee": 1.1954282522201538,
+      "loss_delivery": 0.7644538623946053,
+      "loss_biodist": 0.7812093198299408,
+      "loss_toxic": 0.05465935756053243,
+      "acc_pdi": 0.7076923076923077,
+      "acc_ee": 0.35384615384615387,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.707563672746931,
+      "loss_size": 0.2168926394411496,
+      "loss_pdi": 0.6468359615121569,
+      "loss_ee": 1.2361225570951189,
+      "loss_delivery": 0.8388645563806806,
+      "loss_biodist": 0.7232611009052822,
+      "loss_toxic": 0.04558686592749187,
+      "acc_pdi": 0.7076923076923077,
+      "acc_ee": 0.37435897435897436,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.5122547830854143,
+      "loss_size": 0.2614529473440988,
+      "loss_pdi": 0.6344352598701205,
+      "loss_ee": 1.2337199449539185,
+      "loss_delivery": 0.7557642417294639,
+      "loss_biodist": 0.5974612619195666,
+      "loss_toxic": 0.029421218537858555,
+      "acc_pdi": 0.7076923076923077,
+      "acc_ee": 0.38461538461538464,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 4.085699932915824,
+      "loss_size": 0.2200212436062949,
+      "loss_pdi": 0.6225023801837649,
+      "loss_ee": 1.1934180770601546,
+      "loss_delivery": 1.4556497505732946,
+      "loss_biodist": 0.5685178296906608,
+      "loss_toxic": 0.02559054403432778,
       "acc_pdi": 0.7076923076923077,
       "acc_ee": 0.4205128205128205,
       "acc_toxic": 1.0
     },
     {
-      "loss": 4.04153094972883,
-      "loss_size": 0.18925889155694417,
-      "loss_pdi": 0.6154808104038239,
-      "loss_ee": 1.4497678790773665,
-      "loss_delivery": 0.6580333965165275,
-      "loss_biodist": 1.0802616902760096,
-      "loss_toxic": 0.04872834203498704,
-      "acc_pdi": 0.7076923076923077,
+      "loss": 3.8555148329053606,
+      "loss_size": 0.25200862703578814,
+      "loss_pdi": 0.6161004496472222,
+      "loss_ee": 1.2155306509562902,
+      "loss_delivery": 1.1912458751882826,
+      "loss_biodist": 0.5578728743961879,
+      "loss_toxic": 0.022756420208939483,
+      "acc_pdi": 0.7128205128205128,
+      "acc_ee": 0.4153846153846154,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 4.016897848674229,
+      "loss_size": 0.24823468178510666,
+      "loss_pdi": 0.6098369508981705,
+      "loss_ee": 1.2048260143824987,
+      "loss_delivery": 1.3509162408964974,
+      "loss_biodist": 0.5822246244975499,
+      "loss_toxic": 0.020859350051198686,
+      "acc_pdi": 0.717948717948718,
+      "acc_ee": 0.4153846153846154,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.9899745328085765,
+      "loss_size": 0.27859273659331457,
+      "loss_pdi": 0.6051683936800275,
+      "loss_ee": 1.1875721216201782,
+      "loss_delivery": 1.369072552238192,
+      "loss_biodist": 0.5321473862443652,
+      "loss_toxic": 0.017421354805784568,
+      "acc_pdi": 0.717948717948718,
+      "acc_ee": 0.40512820512820513,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 4.112551552908761,
+      "loss_size": 0.23502862347023828,
+      "loss_pdi": 0.6127274334430695,
+      "loss_ee": 1.2102909428732735,
+      "loss_delivery": 1.5001615626471383,
+      "loss_biodist": 0.5411617543016162,
+      "loss_toxic": 0.013181165459432773,
+      "acc_pdi": 0.7025641025641025,
       "acc_ee": 0.4205128205128205,
       "acc_toxic": 1.0
     },
     {
-      "loss": 4.0117951801845,
-      "loss_size": 0.261259101331234,
-      "loss_pdi": 0.6064834722450801,
-      "loss_ee": 1.4382908344268799,
-      "loss_delivery": 0.6485261533941541,
-      "loss_biodist": 1.013489259140832,
-      "loss_toxic": 0.04374648204871586,
-      "acc_pdi": 0.7076923076923077,
-      "acc_ee": 0.4205128205128205,
+      "loss": 4.031796966280256,
+      "loss_size": 0.27996559121779035,
+      "loss_pdi": 0.6218498295971325,
+      "loss_ee": 1.264663577079773,
+      "loss_delivery": 1.225053642477308,
+      "loss_biodist": 0.6304309921605247,
+      "loss_toxic": 0.009833223053387232,
+      "acc_pdi": 0.7128205128205128,
+      "acc_ee": 0.4153846153846154,
       "acc_toxic": 1.0
     },
     {
-      "loss": 4.021411214556013,
-      "loss_size": 0.25430602048124584,
-      "loss_pdi": 0.6047579922846386,
-      "loss_ee": 1.4375121252877372,
-      "loss_delivery": 0.794183360678809,
-      "loss_biodist": 0.8880592542035239,
-      "loss_toxic": 0.04259242117404938,
-      "acc_pdi": 0.7076923076923077,
-      "acc_ee": 0.41025641025641024,
+      "loss": 4.108100175857544,
+      "loss_size": 0.2646343271647181,
+      "loss_pdi": 0.6244613996573857,
+      "loss_ee": 1.2785721336092268,
+      "loss_delivery": 1.2817653375012534,
+      "loss_biodist": 0.6491539776325226,
+      "loss_toxic": 0.009512946475297213,
+      "acc_pdi": 0.7025641025641025,
+      "acc_ee": 0.4,
       "acc_toxic": 1.0
     },
     {
-      "loss": 4.389863218579974,
-      "loss_size": 0.47541433785642895,
-      "loss_pdi": 0.6053804946797234,
-      "loss_ee": 1.5225199971880232,
-      "loss_delivery": 0.864760024206979,
-      "loss_biodist": 0.8813395031860897,
-      "loss_toxic": 0.04044891867254462,
-      "acc_pdi": 0.7076923076923077,
-      "acc_ee": 0.4205128205128205,
+      "loss": 4.237571648188999,
+      "loss_size": 0.24608339369297028,
+      "loss_pdi": 0.6387959729347911,
+      "loss_ee": 1.263997495174408,
+      "loss_delivery": 1.3677963316440582,
+      "loss_biodist": 0.7113959235804421,
+      "loss_toxic": 0.00950257752888969,
+      "acc_pdi": 0.7025641025641025,
+      "acc_ee": 0.40512820512820513,
       "acc_toxic": 1.0
     },
     {
-      "loss": 4.2736298356737406,
-      "loss_size": 0.4230478884918349,
-      "loss_pdi": 0.6034957447222301,
-      "loss_ee": 1.4814096518925257,
-      "loss_delivery": 0.9466098589556557,
-      "loss_biodist": 0.7808983538831983,
-      "loss_toxic": 0.03816834863807474,
+      "loss": 4.677373443331037,
+      "loss_size": 0.25485377439430784,
+      "loss_pdi": 0.6646602579525539,
+      "loss_ee": 1.2957249539239066,
+      "loss_delivery": 1.609152581010546,
+      "loss_biodist": 0.8443670613425118,
+      "loss_toxic": 0.008614770535911833,
+      "acc_pdi": 0.7025641025641025,
+      "acc_ee": 0.4,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 4.661478791918073,
+      "loss_size": 0.22112409876925604,
+      "loss_pdi": 0.6840873499001775,
+      "loss_ee": 1.345719371523176,
+      "loss_delivery": 1.5750049437795366,
+      "loss_biodist": 0.8278610365731376,
+      "loss_toxic": 0.0076820029810603175,
+      "acc_pdi": 0.6974358974358974,
+      "acc_ee": 0.39487179487179486,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 4.7805344717843195,
+      "loss_size": 0.23743291412081038,
+      "loss_pdi": 0.6911627639617238,
+      "loss_ee": 1.3796877009528024,
+      "loss_delivery": 1.6191245743206568,
+      "loss_biodist": 0.8459444258894239,
+      "loss_toxic": 0.00718201809961881,
       "acc_pdi": 0.7076923076923077,
       "acc_ee": 0.4,
       "acc_toxic": 1.0
     },
     {
-      "loss": 4.585977554321289,
-      "loss_size": 0.4787849709391594,
-      "loss_pdi": 0.6088685957448823,
-      "loss_ee": 1.5371792827333723,
-      "loss_delivery": 1.1909210256167821,
-      "loss_biodist": 0.7510551980563572,
-      "loss_toxic": 0.01916840286659343,
-      "acc_pdi": 0.7076923076923077,
-      "acc_ee": 0.4205128205128205,
+      "loss": 5.050536905016218,
+      "loss_size": 0.2340430138366563,
+      "loss_pdi": 0.7009050281984466,
+      "loss_ee": 1.3757387569972448,
+      "loss_delivery": 1.760018629687173,
+      "loss_biodist": 0.9723425933292934,
+      "loss_toxic": 0.007489000846232686,
+      "acc_pdi": 0.7128205128205128,
+      "acc_ee": 0.39487179487179486,
       "acc_toxic": 1.0
     },
     {
-      "loss": 4.317904029573713,
-      "loss_size": 0.46495475407157627,
-      "loss_pdi": 0.6024806403688022,
-      "loss_ee": 1.5152796847479684,
-      "loss_delivery": 1.0274277882916587,
-      "loss_biodist": 0.6894632875919342,
-      "loss_toxic": 0.018297837914100716,
-      "acc_pdi": 0.7076923076923077,
-      "acc_ee": 0.4205128205128205,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 4.287515469959804,
-      "loss_size": 0.5376717469521931,
-      "loss_pdi": 0.5981708924685206,
-      "loss_ee": 1.4651154450007848,
-      "loss_delivery": 1.0260179340839386,
-      "loss_biodist": 0.6434041815144675,
-      "loss_toxic": 0.01713526714593172,
-      "acc_pdi": 0.7076923076923077,
+      "loss": 5.172625916344779,
+      "loss_size": 0.23549717558281763,
+      "loss_pdi": 0.6980976568801063,
+      "loss_ee": 1.357451047216143,
+      "loss_delivery": 1.893591480595725,
+      "loss_biodist": 0.9802024279321943,
+      "loss_toxic": 0.0077861944612647805,
+      "acc_pdi": 0.717948717948718,
       "acc_ee": 0.38974358974358975,
       "acc_toxic": 1.0
     },
     {
-      "loss": 4.683320045471191,
-      "loss_size": 0.5271316319704056,
-      "loss_pdi": 0.6275609007903508,
-      "loss_ee": 1.5593642677579607,
-      "loss_delivery": 1.2529506555625372,
-      "loss_biodist": 0.7023919905935015,
-      "loss_toxic": 0.013920623875622238,
-      "acc_pdi": 0.7076923076923077,
-      "acc_ee": 0.4205128205128205,
+      "loss": 5.048826490129743,
+      "loss_size": 0.2420537450483867,
+      "loss_pdi": 0.7013733184763363,
+      "loss_ee": 1.353548560823713,
+      "loss_delivery": 1.7931698901312692,
+      "loss_biodist": 0.9509889696325574,
+      "loss_toxic": 0.007692053714501006,
+      "acc_pdi": 0.7128205128205128,
+      "acc_ee": 0.38974358974358975,
       "acc_toxic": 1.0
     },
     {
-      "loss": 4.397426230566842,
-      "loss_size": 0.5106346564633506,
-      "loss_pdi": 0.6181366369128227,
-      "loss_ee": 1.5676180635179793,
-      "loss_delivery": 1.0354454517364502,
-      "loss_biodist": 0.6522124835423061,
-      "loss_toxic": 0.013378978440804141,
+      "loss": 4.951304980686733,
+      "loss_size": 0.24394649054322923,
+      "loss_pdi": 0.7197230202811105,
+      "loss_ee": 1.3789095027106149,
+      "loss_delivery": 1.6561105762209212,
+      "loss_biodist": 0.9460095167160034,
+      "loss_toxic": 0.006605847106714334,
       "acc_pdi": 0.7076923076923077,
-      "acc_ee": 0.41025641025641024,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 4.412310838699341,
-      "loss_size": 0.5362906115395683,
-      "loss_pdi": 0.6130111419728824,
-      "loss_ee": 1.5929286650248937,
-      "loss_delivery": 1.001524874142238,
-      "loss_biodist": 0.6563895855631147,
-      "loss_toxic": 0.012165952806494065,
-      "acc_pdi": 0.7076923076923077,
-      "acc_ee": 0.4205128205128205,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 4.78386572429112,
-      "loss_size": 0.5134438020842416,
-      "loss_pdi": 0.6321895952735629,
-      "loss_ee": 1.5953963484082903,
-      "loss_delivery": 1.3276494145393372,
-      "loss_biodist": 0.7046042127268655,
-      "loss_toxic": 0.010582369419613056,
-      "acc_pdi": 0.7076923076923077,
-      "acc_ee": 0.4205128205128205,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 4.835996287209647,
-      "loss_size": 0.4968718098742621,
-      "loss_pdi": 0.6234481834939548,
-      "loss_ee": 1.5689009257725306,
-      "loss_delivery": 1.4453607542174203,
-      "loss_biodist": 0.6902376328195844,
-      "loss_toxic": 0.011176949766065394,
-      "acc_pdi": 0.7076923076923077,
-      "acc_ee": 0.41025641025641024,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 4.646918603352138,
-      "loss_size": 0.49231292733124327,
-      "loss_pdi": 0.6135136048708644,
-      "loss_ee": 1.5504994562694006,
-      "loss_delivery": 1.3212553603308541,
-      "loss_biodist": 0.6578856153147561,
-      "loss_toxic": 0.01145164788301502,
-      "acc_pdi": 0.7076923076923077,
-      "acc_ee": 0.40512820512820513,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 4.580311230250767,
-      "loss_size": 0.529256243790899,
-      "loss_pdi": 0.6140108300106866,
-      "loss_ee": 1.5573444025857108,
-      "loss_delivery": 1.2449579962662287,
-      "loss_biodist": 0.6246597383703504,
-      "loss_toxic": 0.01008199481293559,
-      "acc_pdi": 0.7076923076923077,
-      "acc_ee": 0.41025641025641024,
+      "acc_ee": 0.38974358974358975,
       "acc_toxic": 1.0
     }
   ]
diff --git a/models/finetune_cv/fold_2/model.pt b/models/finetune_cv/fold_2/model.pt
index 8b04371..d8f6fa3 100644
Binary files a/models/finetune_cv/fold_2/model.pt and b/models/finetune_cv/fold_2/model.pt differ
diff --git a/models/finetune_cv/fold_3/history.json b/models/finetune_cv/fold_3/history.json
index 514aa3e..c35b621 100644
--- a/models/finetune_cv/fold_3/history.json
+++ b/models/finetune_cv/fold_3/history.json
@@ -1,572 +1,446 @@
 {
   "train": [
     {
-      "loss": 21.41742353439331,
-      "loss_size": 15.701433181762695,
-      "loss_pdi": 1.3632826924324035,
-      "loss_ee": 1.0266466200351716,
-      "loss_delivery": 1.418760311603546,
-      "loss_biodist": 1.2302423000335694,
-      "loss_toxic": 0.6770588457584381
+      "loss": 19.317236709594727,
+      "loss_size": 14.108779287338256,
+      "loss_pdi": 1.2223044037818909,
+      "loss_ee": 1.1201724767684937,
+      "loss_delivery": 1.0989456713199615,
+      "loss_biodist": 1.2243955612182618,
+      "loss_toxic": 0.5426396489143371
     },
     {
-      "loss": 9.57167477607727,
-      "loss_size": 4.903112936019897,
-      "loss_pdi": 1.140204393863678,
-      "loss_ee": 1.0074746072292329,
-      "loss_delivery": 1.1362140402197838,
-      "loss_biodist": 1.0133467674255372,
-      "loss_toxic": 0.3713219165802002
+      "loss": 7.058736562728882,
+      "loss_size": 2.407193088531494,
+      "loss_pdi": 1.0396445631980895,
+      "loss_ee": 1.043432891368866,
+      "loss_delivery": 1.183875671029091,
+      "loss_biodist": 1.0435532987117768,
+      "loss_toxic": 0.3410369783639908
     },
     {
-      "loss": 5.485435938835144,
-      "loss_size": 1.4459647357463836,
-      "loss_pdi": 0.9086057424545289,
-      "loss_ee": 0.9852841377258301,
-      "loss_delivery": 1.0618216753005982,
-      "loss_biodist": 0.8613051772117615,
-      "loss_toxic": 0.22245449870824813
+      "loss": 4.30960967540741,
+      "loss_size": 0.3954987198114395,
+      "loss_pdi": 0.8135693371295929,
+      "loss_ee": 0.9796469569206238,
+      "loss_delivery": 1.0894740536808967,
+      "loss_biodist": 0.8407172799110413,
+      "loss_toxic": 0.1907033085823059
     },
     {
-      "loss": 4.065439391136169,
-      "loss_size": 0.45657781660556795,
-      "loss_pdi": 0.7768359065055848,
-      "loss_ee": 0.953633052110672,
-      "loss_delivery": 0.9996832102537155,
-      "loss_biodist": 0.7116856783628464,
-      "loss_toxic": 0.16702365390956403
+      "loss": 3.814995551109314,
+      "loss_size": 0.3033123552799225,
+      "loss_pdi": 0.7126355469226837,
+      "loss_ee": 0.9527663111686706,
+      "loss_delivery": 1.0963637247681617,
+      "loss_biodist": 0.613499540090561,
+      "loss_toxic": 0.13641793020069598
     },
     {
-      "loss": 3.687756299972534,
-      "loss_size": 0.30417054146528244,
-      "loss_pdi": 0.7158299326896668,
-      "loss_ee": 0.9035742998123169,
-      "loss_delivery": 1.060212180018425,
-      "loss_biodist": 0.5556510210037231,
-      "loss_toxic": 0.1483182568103075
+      "loss": 3.4925455331802366,
+      "loss_size": 0.31617238447070123,
+      "loss_pdi": 0.6569374144077301,
+      "loss_ee": 0.9212668299674988,
+      "loss_delivery": 1.0250366538763047,
+      "loss_biodist": 0.4527473896741867,
+      "loss_toxic": 0.12038486637175083
     },
     {
-      "loss": 3.536502242088318,
-      "loss_size": 0.3545849896967411,
-      "loss_pdi": 0.6672206580638885,
-      "loss_ee": 0.8711013793945312,
-      "loss_delivery": 1.0780200093984604,
-      "loss_biodist": 0.4400379478931427,
-      "loss_toxic": 0.12553719095885754
+      "loss": 3.255272912979126,
+      "loss_size": 0.27022836059331895,
+      "loss_pdi": 0.6289663434028625,
+      "loss_ee": 0.9047561466693879,
+      "loss_delivery": 0.9742588266730309,
+      "loss_biodist": 0.3838836058974266,
+      "loss_toxic": 0.09317965283989907
     },
     {
-      "loss": 3.289654517173767,
-      "loss_size": 0.3215350516140461,
-      "loss_pdi": 0.6445666253566742,
-      "loss_ee": 0.8839016199111939,
-      "loss_delivery": 0.9652104169130326,
-      "loss_biodist": 0.35911705493927004,
-      "loss_toxic": 0.11532373651862145
+      "loss": 3.281973719596863,
+      "loss_size": 0.26578598394989966,
+      "loss_pdi": 0.5881433010101318,
+      "loss_ee": 0.8700660288333892,
+      "loss_delivery": 1.1519657298922539,
+      "loss_biodist": 0.3240764126181602,
+      "loss_toxic": 0.08193630240857601
     },
     {
-      "loss": 3.199695384502411,
-      "loss_size": 0.29793725311756136,
-      "loss_pdi": 0.6349938422441482,
-      "loss_ee": 0.8508742034435273,
-      "loss_delivery": 0.9878654226660728,
-      "loss_biodist": 0.3253925606608391,
-      "loss_toxic": 0.10263208523392678
+      "loss": 2.7810576915740968,
+      "loss_size": 0.2505718767642975,
+      "loss_pdi": 0.5593925356864929,
+      "loss_ee": 0.8196510195732116,
+      "loss_delivery": 0.7932070523500443,
+      "loss_biodist": 0.27653754949569703,
+      "loss_toxic": 0.08169759791344404
     },
     {
-      "loss": 2.981464409828186,
-      "loss_size": 0.2926298946142197,
-      "loss_pdi": 0.6041542202234268,
-      "loss_ee": 0.8202637135982513,
-      "loss_delivery": 0.9215006068348884,
-      "loss_biodist": 0.2569382354617119,
-      "loss_toxic": 0.08597766645252705
+      "loss": 2.644732141494751,
+      "loss_size": 0.27979295402765275,
+      "loss_pdi": 0.5457461476325989,
+      "loss_ee": 0.7845215618610382,
+      "loss_delivery": 0.722122372686863,
+      "loss_biodist": 0.2437703028321266,
+      "loss_toxic": 0.06877880096435547
     },
     {
-      "loss": 2.7703840017318724,
-      "loss_size": 0.30560431331396104,
-      "loss_pdi": 0.5837061107158661,
-      "loss_ee": 0.7886367738246918,
-      "loss_delivery": 0.763620425760746,
-      "loss_biodist": 0.23922762870788575,
-      "loss_toxic": 0.08958881739526987
+      "loss": 2.5743841886520387,
+      "loss_size": 0.21236803606152535,
+      "loss_pdi": 0.5281321376562118,
+      "loss_ee": 0.7772053182125092,
+      "loss_delivery": 0.7842913195490837,
+      "loss_biodist": 0.20931598618626596,
+      "loss_toxic": 0.0630713876336813
     },
     {
-      "loss": 2.6885447978973387,
-      "loss_size": 0.28905282765626905,
-      "loss_pdi": 0.5664507508277893,
-      "loss_ee": 0.7639070689678192,
-      "loss_delivery": 0.7704478114843368,
-      "loss_biodist": 0.21979134678840637,
-      "loss_toxic": 0.07889490202069283
+      "loss": 2.493379771709442,
+      "loss_size": 0.2545281477272511,
+      "loss_pdi": 0.514763566851616,
+      "loss_ee": 0.7416582465171814,
+      "loss_delivery": 0.7315813854336739,
+      "loss_biodist": 0.18844463676214218,
+      "loss_toxic": 0.062403830140829085
     },
     {
-      "loss": 2.591668117046356,
-      "loss_size": 0.2613472960889339,
-      "loss_pdi": 0.5540884166955948,
-      "loss_ee": 0.7465697586536407,
-      "loss_delivery": 0.7534924671053886,
-      "loss_biodist": 0.20271824076771736,
-      "loss_toxic": 0.07345191687345505
+      "loss": 2.3714203119277952,
+      "loss_size": 0.21288565024733544,
+      "loss_pdi": 0.5149440914392471,
+      "loss_ee": 0.7432775914669036,
+      "loss_delivery": 0.6615208894014358,
+      "loss_biodist": 0.17799324095249175,
+      "loss_toxic": 0.06079882858321071
     },
     {
-      "loss": 2.482152557373047,
-      "loss_size": 0.2550207316875458,
-      "loss_pdi": 0.5377364099025727,
-      "loss_ee": 0.7093640863895416,
-      "loss_delivery": 0.715790644288063,
-      "loss_biodist": 0.1968037411570549,
-      "loss_toxic": 0.0674369728192687
+      "loss": 2.3138927936553957,
+      "loss_size": 0.22406778559088708,
+      "loss_pdi": 0.5060430943965912,
+      "loss_ee": 0.7270951688289642,
+      "loss_delivery": 0.6268678307533264,
+      "loss_biodist": 0.17946239709854125,
+      "loss_toxic": 0.050356499617919326
     },
     {
-      "loss": 2.4853516697883604,
-      "loss_size": 0.25445577800273894,
-      "loss_pdi": 0.5285622417926789,
-      "loss_ee": 0.7045736134052276,
-      "loss_delivery": 0.765293450653553,
-      "loss_biodist": 0.16704678535461426,
-      "loss_toxic": 0.06541981641203165
+      "loss": 2.2404407501220702,
+      "loss_size": 0.23460092321038245,
+      "loss_pdi": 0.4892877459526062,
+      "loss_ee": 0.6908941030502319,
+      "loss_delivery": 0.6124202072620392,
+      "loss_biodist": 0.16842604279518128,
+      "loss_toxic": 0.044811736792325974
     },
     {
-      "loss": 2.325911021232605,
-      "loss_size": 0.22311154529452323,
-      "loss_pdi": 0.5131498754024506,
-      "loss_ee": 0.6778043508529663,
-      "loss_delivery": 0.6851547978818416,
-      "loss_biodist": 0.16465196907520294,
-      "loss_toxic": 0.06203848272562027
+      "loss": 2.2448294520378114,
+      "loss_size": 0.21119624376296997,
+      "loss_pdi": 0.479864901304245,
+      "loss_ee": 0.6906192302703857,
+      "loss_delivery": 0.6555144399404526,
+      "loss_biodist": 0.16310803219676018,
+      "loss_toxic": 0.044526621932163835
     },
     {
-      "loss": 2.213776695728302,
-      "loss_size": 0.24956582188606263,
-      "loss_pdi": 0.48536362051963805,
-      "loss_ee": 0.6791646689176559,
-      "loss_delivery": 0.5861033886671067,
-      "loss_biodist": 0.15450086519122125,
-      "loss_toxic": 0.059078306704759595
+      "loss": 2.1580574989318846,
+      "loss_size": 0.18697498068213464,
+      "loss_pdi": 0.48660930395126345,
+      "loss_ee": 0.6810935467481614,
+      "loss_delivery": 0.6051739566028118,
+      "loss_biodist": 0.15406969040632248,
+      "loss_toxic": 0.04413598729297519
     },
     {
-      "loss": 2.3086095094680785,
-      "loss_size": 0.20203103460371494,
-      "loss_pdi": 0.4988987535238266,
-      "loss_ee": 0.6785910665988922,
-      "loss_delivery": 0.7291694968938828,
-      "loss_biodist": 0.1462075024843216,
-      "loss_toxic": 0.05371163971722126
+      "loss": 2.114891529083252,
+      "loss_size": 0.17799586579203605,
+      "loss_pdi": 0.4589719235897064,
+      "loss_ee": 0.6686563313007354,
+      "loss_delivery": 0.6179293170571327,
+      "loss_biodist": 0.1526280902326107,
+      "loss_toxic": 0.03870999766513705
     },
     {
-      "loss": 2.0882336378097532,
-      "loss_size": 0.2287739872932434,
-      "loss_pdi": 0.5072675496339798,
-      "loss_ee": 0.6701794564723969,
-      "loss_delivery": 0.4832353606820107,
-      "loss_biodist": 0.1421804867684841,
-      "loss_toxic": 0.05659680655226111
+      "loss": 2.1680126667022703,
+      "loss_size": 0.18272313922643663,
+      "loss_pdi": 0.47693236321210863,
+      "loss_ee": 0.6723115026950837,
+      "loss_delivery": 0.6574018053710461,
+      "loss_biodist": 0.14306045994162558,
+      "loss_toxic": 0.03558342705946416
     },
     {
-      "loss": 2.0332674741744996,
-      "loss_size": 0.22676872164011003,
-      "loss_pdi": 0.4669553279876709,
-      "loss_ee": 0.6482236534357071,
-      "loss_delivery": 0.5079896807670593,
-      "loss_biodist": 0.13781072497367858,
-      "loss_toxic": 0.04551935400813818
+      "loss": 2.0243090748786927,
+      "loss_size": 0.19451010078191758,
+      "loss_pdi": 0.46296934187412264,
+      "loss_ee": 0.6654580652713775,
+      "loss_delivery": 0.5195972554385662,
+      "loss_biodist": 0.1416195034980774,
+      "loss_toxic": 0.04015482016839087
     },
     {
-      "loss": 1.9842296838760376,
-      "loss_size": 0.2007827118039131,
-      "loss_pdi": 0.4655294865369797,
-      "loss_ee": 0.6295293152332306,
-      "loss_delivery": 0.5032630048692226,
-      "loss_biodist": 0.12692373394966125,
-      "loss_toxic": 0.05820149295032025
+      "loss": 1.980038857460022,
+      "loss_size": 0.19992023780941964,
+      "loss_pdi": 0.4373833805322647,
+      "loss_ee": 0.6562270969152451,
+      "loss_delivery": 0.5170416861772538,
+      "loss_biodist": 0.13248837292194365,
+      "loss_toxic": 0.03697808152064681
     },
     {
-      "loss": 1.9574703454971314,
-      "loss_size": 0.20575413331389428,
-      "loss_pdi": 0.4662397414445877,
-      "loss_ee": 0.6207350552082062,
-      "loss_delivery": 0.4868774816393852,
-      "loss_biodist": 0.1334032252430916,
-      "loss_toxic": 0.04446075968444348
-    },
-    {
-      "loss": 1.8693695425987245,
-      "loss_size": 0.2003278151154518,
-      "loss_pdi": 0.4481669098138809,
-      "loss_ee": 0.6228156566619873,
-      "loss_delivery": 0.4377666234970093,
-      "loss_biodist": 0.1222815040498972,
-      "loss_toxic": 0.03801098903641105
-    },
-    {
-      "loss": 1.9101393103599549,
-      "loss_size": 0.22527543231844901,
-      "loss_pdi": 0.4501156389713287,
-      "loss_ee": 0.5992490768432617,
-      "loss_delivery": 0.4722588837146759,
-      "loss_biodist": 0.12016024515032768,
-      "loss_toxic": 0.04308005180209875
-    },
-    {
-      "loss": 1.8186616897583008,
-      "loss_size": 0.20228504091501237,
-      "loss_pdi": 0.4353774756193161,
-      "loss_ee": 0.595091101527214,
-      "loss_delivery": 0.42390944324433805,
-      "loss_biodist": 0.12379681393504142,
-      "loss_toxic": 0.03820184739306569
-    },
-    {
-      "loss": 1.753528320789337,
-      "loss_size": 0.168972497433424,
-      "loss_pdi": 0.4384701639413834,
-      "loss_ee": 0.5804536670446396,
-      "loss_delivery": 0.4066555552184582,
-      "loss_biodist": 0.12041406258940697,
-      "loss_toxic": 0.038562366552650926
-    },
-    {
-      "loss": 1.752294898033142,
-      "loss_size": 0.16247735619544984,
-      "loss_pdi": 0.4280344098806381,
-      "loss_ee": 0.5781133621931076,
-      "loss_delivery": 0.4378485083580017,
-      "loss_biodist": 0.1110315527766943,
-      "loss_toxic": 0.03478973265737295
-    },
-    {
-      "loss": 1.8295514345169068,
-      "loss_size": 0.1906499370932579,
-      "loss_pdi": 0.42695977091789244,
-      "loss_ee": 0.6036634147167206,
-      "loss_delivery": 0.4585780970752239,
-      "loss_biodist": 0.10937272906303405,
-      "loss_toxic": 0.04032747393939644
+      "loss": 2.0073827385902403,
+      "loss_size": 0.17545675858855247,
+      "loss_pdi": 0.43559625148773196,
+      "loss_ee": 0.6394164443016053,
+      "loss_delivery": 0.5809337809681893,
+      "loss_biodist": 0.13504885137081146,
+      "loss_toxic": 0.040930699557065964
     }
   ],
   "val": [
     {
-      "loss": 12.90866756439209,
-      "loss_size": 8.614875793457031,
-      "loss_pdi": 1.3393473625183105,
-      "loss_ee": 0.7915782928466797,
-      "loss_delivery": 0.35056664049625397,
-      "loss_biodist": 1.2413995265960693,
-      "loss_toxic": 0.5708996057510376,
-      "acc_pdi": 0.09803921568627451,
-      "acc_ee": 0.8431372549019608,
+      "loss": 10.945204257965088,
+      "loss_size": 6.681218147277832,
+      "loss_pdi": 1.0216107964515686,
+      "loss_ee": 1.0486068725585938,
+      "loss_delivery": 0.4687899202108383,
+      "loss_biodist": 1.215298354625702,
+      "loss_toxic": 0.5096809715032578,
+      "acc_pdi": 0.8823529411764706,
+      "acc_ee": 0.6470588235294118,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 6.046682596206665,
-      "loss_size": 2.181392192840576,
-      "loss_pdi": 1.010448306798935,
-      "loss_ee": 0.7943653464317322,
-      "loss_delivery": 0.3778253495693207,
-      "loss_biodist": 1.1818514466285706,
-      "loss_toxic": 0.5007997751235962,
+      "loss": 4.0456085205078125,
+      "loss_size": 0.3470493406057358,
+      "loss_pdi": 0.7843169867992401,
+      "loss_ee": 0.8336820006370544,
+      "loss_delivery": 0.42519159615039825,
+      "loss_biodist": 1.1528617143630981,
+      "loss_toxic": 0.5025068372488022,
       "acc_pdi": 0.8823529411764706,
       "acc_ee": 0.8431372549019608,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 4.130645751953125,
-      "loss_size": 0.4065478593111038,
-      "loss_pdi": 0.7734719514846802,
-      "loss_ee": 0.7856209874153137,
-      "loss_delivery": 0.4249718487262726,
-      "loss_biodist": 1.1932182908058167,
-      "loss_toxic": 0.5468149781227112,
+      "loss": 3.553251266479492,
+      "loss_size": 0.07565776817500591,
+      "loss_pdi": 0.5630811750888824,
+      "loss_ee": 0.6947644650936127,
+      "loss_delivery": 0.4020952582359314,
+      "loss_biodist": 1.1898333430290222,
+      "loss_toxic": 0.627819336950779,
       "acc_pdi": 0.8823529411764706,
       "acc_ee": 0.8431372549019608,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.74097740650177,
-      "loss_size": 0.05154523253440857,
-      "loss_pdi": 0.6415763646364212,
-      "loss_ee": 0.7410732209682465,
-      "loss_delivery": 0.38766030967235565,
-      "loss_biodist": 1.2818186283111572,
-      "loss_toxic": 0.6373037025332451,
+      "loss": 3.2876256704330444,
+      "loss_size": 0.056399866938591,
+      "loss_pdi": 0.579200953245163,
+      "loss_ee": 0.5947848558425903,
+      "loss_delivery": 0.4561047703027725,
+      "loss_biodist": 1.0357274413108826,
+      "loss_toxic": 0.5654077678918839,
       "acc_pdi": 0.8823529411764706,
       "acc_ee": 0.8431372549019608,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.6484163999557495,
-      "loss_size": 0.05830332264304161,
-      "loss_pdi": 0.5795184522867203,
-      "loss_ee": 0.6541054546833038,
-      "loss_delivery": 0.40608713030815125,
-      "loss_biodist": 1.2502482235431671,
-      "loss_toxic": 0.7001538649201393,
+      "loss": 3.2625988721847534,
+      "loss_size": 0.11296019703149796,
+      "loss_pdi": 0.5352367609739304,
+      "loss_ee": 0.6021667718887329,
+      "loss_delivery": 0.5046610683202744,
+      "loss_biodist": 1.0080225467681885,
+      "loss_toxic": 0.4995514266192913,
       "acc_pdi": 0.8823529411764706,
       "acc_ee": 0.8431372549019608,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.8285857439041138,
-      "loss_size": 0.07835924997925758,
-      "loss_pdi": 0.5752626657485962,
-      "loss_ee": 0.7419937252998352,
-      "loss_delivery": 0.382533997297287,
-      "loss_biodist": 1.3310475945472717,
-      "loss_toxic": 0.7193883210420609,
+      "loss": 3.165306806564331,
+      "loss_size": 0.09038551151752472,
+      "loss_pdi": 0.5058617442846298,
+      "loss_ee": 0.6476156711578369,
+      "loss_delivery": 0.43027013540267944,
+      "loss_biodist": 1.0445645153522491,
+      "loss_toxic": 0.4466092698276043,
       "acc_pdi": 0.8823529411764706,
       "acc_ee": 0.8431372549019608,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.6414456367492676,
-      "loss_size": 0.07554847374558449,
-      "loss_pdi": 0.6094752848148346,
-      "loss_ee": 0.7212981283664703,
-      "loss_delivery": 0.4330318123102188,
-      "loss_biodist": 1.1887046694755554,
-      "loss_toxic": 0.6133871823549271,
+      "loss": 3.2408690452575684,
+      "loss_size": 0.22243183851242065,
+      "loss_pdi": 0.4985402673482895,
+      "loss_ee": 0.571824312210083,
+      "loss_delivery": 0.43825456500053406,
+      "loss_biodist": 0.9937507510185242,
+      "loss_toxic": 0.5160673335194588,
       "acc_pdi": 0.8823529411764706,
       "acc_ee": 0.8431372549019608,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.537736415863037,
-      "loss_size": 0.10301502048969269,
-      "loss_pdi": 0.5892050117254257,
-      "loss_ee": 0.6979120671749115,
-      "loss_delivery": 0.5159651935100555,
-      "loss_biodist": 1.064111590385437,
-      "loss_toxic": 0.5675273537635803,
+      "loss": 3.5194804668426514,
+      "loss_size": 0.36968255043029785,
+      "loss_pdi": 0.4991031885147095,
+      "loss_ee": 0.5797468274831772,
+      "loss_delivery": 0.5644859671592712,
+      "loss_biodist": 0.9723091125488281,
+      "loss_toxic": 0.5341527052223682,
       "acc_pdi": 0.8823529411764706,
       "acc_ee": 0.8431372549019608,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.6027194261550903,
-      "loss_size": 0.16463171318173409,
-      "loss_pdi": 0.5593049973249435,
-      "loss_ee": 0.6619580686092377,
-      "loss_delivery": 0.538649171590805,
-      "loss_biodist": 1.0742176473140717,
-      "loss_toxic": 0.6039578504860401,
+      "loss": 3.539685606956482,
+      "loss_size": 0.38313066959381104,
+      "loss_pdi": 0.528433233499527,
+      "loss_ee": 0.5810057669878006,
+      "loss_delivery": 0.44039086997509,
+      "loss_biodist": 1.0017918348312378,
+      "loss_toxic": 0.6049331650137901,
       "acc_pdi": 0.8823529411764706,
       "acc_ee": 0.8431372549019608,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.762178421020508,
-      "loss_size": 0.22026687115430832,
-      "loss_pdi": 0.5431422591209412,
-      "loss_ee": 0.7102161943912506,
-      "loss_delivery": 0.44816476106643677,
-      "loss_biodist": 1.172889918088913,
-      "loss_toxic": 0.6674983687698841,
-      "acc_pdi": 0.8823529411764706,
+      "loss": 3.9403724670410156,
+      "loss_size": 0.6225972771644592,
+      "loss_pdi": 0.5688649713993073,
+      "loss_ee": 0.6205386221408844,
+      "loss_delivery": 0.6095166206359863,
+      "loss_biodist": 0.9419751763343811,
+      "loss_toxic": 0.576879795640707,
+      "acc_pdi": 0.803921568627451,
       "acc_ee": 0.8431372549019608,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.7515294551849365,
-      "loss_size": 0.2188272960484028,
-      "loss_pdi": 0.5569919049739838,
-      "loss_ee": 0.6500461399555206,
-      "loss_delivery": 0.4802835136651993,
-      "loss_biodist": 1.140358328819275,
-      "loss_toxic": 0.7050221972167492,
-      "acc_pdi": 0.8823529411764706,
+      "loss": 3.8653494119644165,
+      "loss_size": 0.6294703483581543,
+      "loss_pdi": 0.5615053772926331,
+      "loss_ee": 0.6072992980480194,
+      "loss_delivery": 0.47824281454086304,
+      "loss_biodist": 0.964938759803772,
+      "loss_toxic": 0.6238927394151688,
+      "acc_pdi": 0.8431372549019608,
       "acc_ee": 0.8431372549019608,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.47284197807312,
-      "loss_size": 0.19455180317163467,
-      "loss_pdi": 0.5266519337892532,
-      "loss_ee": 0.6175068914890289,
-      "loss_delivery": 0.5526535362005234,
-      "loss_biodist": 0.9517507255077362,
-      "loss_toxic": 0.6297270879149437,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.8431372549019608,
+      "loss": 4.109289169311523,
+      "loss_size": 0.753549188375473,
+      "loss_pdi": 0.6232334971427917,
+      "loss_ee": 0.6752453744411469,
+      "loss_delivery": 0.4541686922311783,
+      "loss_biodist": 1.0001222491264343,
+      "loss_toxic": 0.6029700562357903,
+      "acc_pdi": 0.7450980392156863,
+      "acc_ee": 0.7254901960784313,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.9374399185180664,
-      "loss_size": 0.32401855289936066,
-      "loss_pdi": 0.5553185790777206,
-      "loss_ee": 0.6632668077945709,
-      "loss_delivery": 0.48028646409511566,
-      "loss_biodist": 1.1518925726413727,
-      "loss_toxic": 0.7626568526029587,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.8431372549019608,
+      "loss": 4.084217548370361,
+      "loss_size": 0.6689053475856781,
+      "loss_pdi": 0.5947604179382324,
+      "loss_ee": 0.6819752305746078,
+      "loss_delivery": 0.5174736380577087,
+      "loss_biodist": 0.9870622158050537,
+      "loss_toxic": 0.6340407878160477,
+      "acc_pdi": 0.7843137254901961,
+      "acc_ee": 0.7647058823529411,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.7636899948120117,
-      "loss_size": 0.2452084794640541,
-      "loss_pdi": 0.5189685672521591,
-      "loss_ee": 0.6509725153446198,
-      "loss_delivery": 0.38138364255428314,
-      "loss_biodist": 1.1877531707286835,
-      "loss_toxic": 0.7794036716222763,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.8235294117647058,
+      "loss": 3.762814164161682,
+      "loss_size": 0.5682831406593323,
+      "loss_pdi": 0.5777421444654465,
+      "loss_ee": 0.7156199663877487,
+      "loss_delivery": 0.44971026480197906,
+      "loss_biodist": 0.9156049191951752,
+      "loss_toxic": 0.5358536541461945,
+      "acc_pdi": 0.7450980392156863,
+      "acc_ee": 0.7058823529411765,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.837575674057007,
-      "loss_size": 0.3500683009624481,
-      "loss_pdi": 0.5655115246772766,
-      "loss_ee": 0.6301239728927612,
-      "loss_delivery": 0.47451916337013245,
-      "loss_biodist": 1.069214403629303,
-      "loss_toxic": 0.7481381297111511,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.8431372549019608,
+      "loss": 4.147223711013794,
+      "loss_size": 0.6644828915596008,
+      "loss_pdi": 0.5911359935998917,
+      "loss_ee": 0.713784396648407,
+      "loss_delivery": 0.500703439116478,
+      "loss_biodist": 1.0310384333133698,
+      "loss_toxic": 0.6460786163806915,
+      "acc_pdi": 0.7647058823529411,
+      "acc_ee": 0.7058823529411765,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.834155559539795,
-      "loss_size": 0.27750423550605774,
-      "loss_pdi": 0.513676255941391,
-      "loss_ee": 0.6291456520557404,
-      "loss_delivery": 0.40621335804462433,
-      "loss_biodist": 1.1680629253387451,
-      "loss_toxic": 0.8395530804991722,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.8235294117647058,
+      "loss": 4.157853841781616,
+      "loss_size": 0.7414849102497101,
+      "loss_pdi": 0.630668580532074,
+      "loss_ee": 0.6938402056694031,
+      "loss_delivery": 0.50765261054039,
+      "loss_biodist": 0.9891600012779236,
+      "loss_toxic": 0.5950475558638573,
+      "acc_pdi": 0.7450980392156863,
+      "acc_ee": 0.7450980392156863,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.6893343925476074,
-      "loss_size": 0.34515636414289474,
-      "loss_pdi": 0.5478167533874512,
-      "loss_ee": 0.6345725357532501,
-      "loss_delivery": 0.42669548094272614,
-      "loss_biodist": 1.0210922062397003,
-      "loss_toxic": 0.7140011191368103,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.803921568627451,
+      "loss": 4.2473719120025635,
+      "loss_size": 0.8058429956436157,
+      "loss_pdi": 0.5982940196990967,
+      "loss_ee": 0.7209844589233398,
+      "loss_delivery": 0.5253763496875763,
+      "loss_biodist": 0.973820835351944,
+      "loss_toxic": 0.6230533868074417,
+      "acc_pdi": 0.803921568627451,
+      "acc_ee": 0.6862745098039216,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 4.00758171081543,
-      "loss_size": 0.3066108226776123,
-      "loss_pdi": 0.5529182702302933,
-      "loss_ee": 0.6956813335418701,
-      "loss_delivery": 0.4755028784275055,
-      "loss_biodist": 1.152823954820633,
-      "loss_toxic": 0.8240445479750633,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.803921568627451,
+      "loss": 3.90485680103302,
+      "loss_size": 0.5400884747505188,
+      "loss_pdi": 0.5671974420547485,
+      "loss_ee": 0.7085212767124176,
+      "loss_delivery": 0.5078988373279572,
+      "loss_biodist": 0.9940473735332489,
+      "loss_toxic": 0.5871035009622574,
+      "acc_pdi": 0.803921568627451,
+      "acc_ee": 0.7450980392156863,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.974515676498413,
-      "loss_size": 0.3346070274710655,
-      "loss_pdi": 0.5447590947151184,
-      "loss_ee": 0.6450685858726501,
-      "loss_delivery": 0.4816073626279831,
-      "loss_biodist": 1.1390976309776306,
-      "loss_toxic": 0.8293759152293205,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.803921568627451,
+      "loss": 4.193094968795776,
+      "loss_size": 0.764777421951294,
+      "loss_pdi": 0.5734306275844574,
+      "loss_ee": 0.7070393562316895,
+      "loss_delivery": 0.5335722267627716,
+      "loss_biodist": 1.0060182809829712,
+      "loss_toxic": 0.6082571670413017,
+      "acc_pdi": 0.8235294117647058,
+      "acc_ee": 0.7450980392156863,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.880821943283081,
-      "loss_size": 0.282451294362545,
-      "loss_pdi": 0.5469185262918472,
-      "loss_ee": 0.6417804062366486,
-      "loss_delivery": 0.46756358444690704,
-      "loss_biodist": 1.127053290605545,
-      "loss_toxic": 0.8150547966361046,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.803921568627451,
+      "loss": 4.225732326507568,
+      "loss_size": 0.7807798981666565,
+      "loss_pdi": 0.57969930768013,
+      "loss_ee": 0.7046914398670197,
+      "loss_delivery": 0.5619150400161743,
+      "loss_biodist": 1.0033797025680542,
+      "loss_toxic": 0.5952669233083725,
+      "acc_pdi": 0.7843137254901961,
+      "acc_ee": 0.7450980392156863,
       "acc_toxic": 0.851063829787234
     },
     {
-      "loss": 3.8784565925598145,
-      "loss_size": 0.2880236804485321,
-      "loss_pdi": 0.5311962813138962,
-      "loss_ee": 0.6260144710540771,
-      "loss_delivery": 0.4471806138753891,
-      "loss_biodist": 1.1587725579738617,
-      "loss_toxic": 0.8272688835859299,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.803921568627451,
-      "acc_toxic": 0.851063829787234
-    },
-    {
-      "loss": 3.9609978199005127,
-      "loss_size": 0.3277027904987335,
-      "loss_pdi": 0.5816705077886581,
-      "loss_ee": 0.6382596492767334,
-      "loss_delivery": 0.42079347372055054,
-      "loss_biodist": 1.1735960245132446,
-      "loss_toxic": 0.8189755231142044,
-      "acc_pdi": 0.9019607843137255,
-      "acc_ee": 0.803921568627451,
-      "acc_toxic": 0.851063829787234
-    },
-    {
-      "loss": 3.8955435752868652,
-      "loss_size": 0.28625721484422684,
-      "loss_pdi": 0.5423711687326431,
-      "loss_ee": 0.6477845013141632,
-      "loss_delivery": 0.43728773295879364,
-      "loss_biodist": 1.1806198358535767,
-      "loss_toxic": 0.801223024725914,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.803921568627451,
-      "acc_toxic": 0.851063829787234
-    },
-    {
-      "loss": 4.007414698600769,
-      "loss_size": 0.3294255882501602,
-      "loss_pdi": 0.5471038520336151,
-      "loss_ee": 0.6502591967582703,
-      "loss_delivery": 0.4436507970094681,
-      "loss_biodist": 1.202763706445694,
-      "loss_toxic": 0.8342117220163345,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.803921568627451,
-      "acc_toxic": 0.851063829787234
-    },
-    {
-      "loss": 4.050978183746338,
-      "loss_size": 0.39942242205142975,
-      "loss_pdi": 0.5590354651212692,
-      "loss_ee": 0.6409733295440674,
-      "loss_delivery": 0.45262810587882996,
-      "loss_biodist": 1.1648951172828674,
-      "loss_toxic": 0.8340234756469727,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.803921568627451,
-      "acc_toxic": 0.851063829787234
-    },
-    {
-      "loss": 3.86995267868042,
-      "loss_size": 0.2950260043144226,
-      "loss_pdi": 0.5328812152147293,
-      "loss_ee": 0.6263198852539062,
-      "loss_delivery": 0.45980168879032135,
-      "loss_biodist": 1.1475371420383453,
-      "loss_toxic": 0.808386467397213,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.803921568627451,
-      "acc_toxic": 0.851063829787234
-    },
-    {
-      "loss": 3.8100762367248535,
-      "loss_size": 0.2826509475708008,
-      "loss_pdi": 0.5344351083040237,
-      "loss_ee": 0.6322177648544312,
-      "loss_delivery": 0.4711516797542572,
-      "loss_biodist": 1.1184114813804626,
-      "loss_toxic": 0.7712092772126198,
-      "acc_pdi": 0.8823529411764706,
-      "acc_ee": 0.803921568627451,
+      "loss": 4.122166633605957,
+      "loss_size": 0.727344423532486,
+      "loss_pdi": 0.5855642706155777,
+      "loss_ee": 0.7065156400203705,
+      "loss_delivery": 0.5201583206653595,
+      "loss_biodist": 0.9953437149524689,
+      "loss_toxic": 0.5872401967644691,
+      "acc_pdi": 0.803921568627451,
+      "acc_ee": 0.7450980392156863,
       "acc_toxic": 0.851063829787234
     }
   ]
diff --git a/models/finetune_cv/fold_3/model.pt b/models/finetune_cv/fold_3/model.pt
index a55c2e7..fbb77cd 100644
Binary files a/models/finetune_cv/fold_3/model.pt and b/models/finetune_cv/fold_3/model.pt differ
diff --git a/models/finetune_cv/fold_4/history.json b/models/finetune_cv/fold_4/history.json
index a4cd136..7183333 100644
--- a/models/finetune_cv/fold_4/history.json
+++ b/models/finetune_cv/fold_4/history.json
@@ -1,404 +1,383 @@
 {
   "train": [
     {
-      "loss": 17.82338151064786,
-      "loss_size": 12.698930350216953,
-      "loss_pdi": 1.2952325452457776,
-      "loss_ee": 1.044789281758395,
-      "loss_delivery": 0.894258591261777,
-      "loss_biodist": 1.3406665433536877,
-      "loss_toxic": 0.5495036569508639
+      "loss": 19.51876787705855,
+      "loss_size": 14.430041096427225,
+      "loss_pdi": 1.3163191405209629,
+      "loss_ee": 1.080450177192688,
+      "loss_delivery": 0.8366337662393396,
+      "loss_biodist": 1.2393495603041216,
+      "loss_toxic": 0.6159739873626016
     },
     {
-      "loss": 6.6719701940363105,
-      "loss_size": 2.361687348647551,
-      "loss_pdi": 0.9808245301246643,
-      "loss_ee": 0.9910268241708929,
-      "loss_delivery": 0.8210993910377676,
-      "loss_biodist": 1.184872735630382,
-      "loss_toxic": 0.3324593657797033
+      "loss": 7.450059110468084,
+      "loss_size": 3.0973785898902197,
+      "loss_pdi": 1.0574429847977378,
+      "loss_ee": 0.972686382857236,
+      "loss_delivery": 0.8131051341241057,
+      "loss_biodist": 1.0609121160073713,
+      "loss_toxic": 0.4485339197245511
     },
     {
-      "loss": 4.197543209249323,
-      "loss_size": 0.3305926668373021,
-      "loss_pdi": 0.8021349581805143,
-      "loss_ee": 0.9533808014609597,
-      "loss_delivery": 0.85791384361007,
-      "loss_biodist": 1.0029216788031838,
-      "loss_toxic": 0.2505992312322963
+      "loss": 4.450224074450406,
+      "loss_size": 0.4227097576314753,
+      "loss_pdi": 0.8444447679953142,
+      "loss_ee": 0.9470934163440358,
+      "loss_delivery": 1.0848771997473456,
+      "loss_biodist": 0.8853748061440208,
+      "loss_toxic": 0.26572414352135226
     },
     {
-      "loss": 3.817467537793246,
-      "loss_size": 0.3109576465053992,
-      "loss_pdi": 0.7077692205255682,
-      "loss_ee": 0.9259536266326904,
-      "loss_delivery": 0.8590928573500026,
-      "loss_biodist": 0.7878239534117959,
-      "loss_toxic": 0.22587016597390175
+      "loss": 3.697209119796753,
+      "loss_size": 0.27599087357521057,
+      "loss_pdi": 0.7228363969109275,
+      "loss_ee": 0.9080322655764493,
+      "loss_delivery": 0.8549216186458414,
+      "loss_biodist": 0.7190906730565158,
+      "loss_toxic": 0.21633729677308688
     },
     {
-      "loss": 3.426318342035467,
-      "loss_size": 0.30332070047205145,
-      "loss_pdi": 0.6738776402039961,
-      "loss_ee": 0.8633853034539656,
-      "loss_delivery": 0.810848053206097,
-      "loss_biodist": 0.6030096492984078,
-      "loss_toxic": 0.17187697711316022
+      "loss": 3.5547448938543145,
+      "loss_size": 0.29429094425656577,
+      "loss_pdi": 0.6620089682665738,
+      "loss_ee": 0.8595339439131997,
+      "loss_delivery": 0.9978644847869873,
+      "loss_biodist": 0.5649206800894304,
+      "loss_toxic": 0.17612605541944504
     },
     {
-      "loss": 3.243085037578236,
-      "loss_size": 0.2643032656474547,
-      "loss_pdi": 0.680456202138554,
-      "loss_ee": 0.8441235260529951,
-      "loss_delivery": 0.790125925432552,
-      "loss_biodist": 0.4977728453549472,
-      "loss_toxic": 0.16630326211452484
+      "loss": 3.069189115004106,
+      "loss_size": 0.30328830602494156,
+      "loss_pdi": 0.6169473230838776,
+      "loss_ee": 0.8132463910362937,
+      "loss_delivery": 0.72551099143245,
+      "loss_biodist": 0.46287189017642627,
+      "loss_toxic": 0.14732415906407617
     },
     {
-      "loss": 3.0106391689994116,
-      "loss_size": 0.2923726256598126,
-      "loss_pdi": 0.6462291912599043,
-      "loss_ee": 0.8259676423939791,
-      "loss_delivery": 0.7123563428494063,
-      "loss_biodist": 0.4204860031604767,
-      "loss_toxic": 0.11322732371362773
+      "loss": 3.1349260156804863,
+      "loss_size": 0.27535233172503387,
+      "loss_pdi": 0.5893999202684923,
+      "loss_ee": 0.807813747362657,
+      "loss_delivery": 0.9538742283528502,
+      "loss_biodist": 0.4018077091737227,
+      "loss_toxic": 0.1066781035201116
     },
     {
-      "loss": 2.6575805924155493,
-      "loss_size": 0.23369696668603204,
-      "loss_pdi": 0.5998414809053595,
-      "loss_ee": 0.7827209342609752,
-      "loss_delivery": 0.6251701191067696,
-      "loss_biodist": 0.324639000675895,
-      "loss_toxic": 0.09151209010319276
+      "loss": 2.6963415037501943,
+      "loss_size": 0.276088840582154,
+      "loss_pdi": 0.5490301495248621,
+      "loss_ee": 0.757920276034962,
+      "loss_delivery": 0.6706068068742752,
+      "loss_biodist": 0.34504769336093555,
+      "loss_toxic": 0.09764780781485817
     },
     {
-      "loss": 2.6034729263999243,
-      "loss_size": 0.20161977207118814,
-      "loss_pdi": 0.5999934835867449,
-      "loss_ee": 0.7699762636964972,
-      "loss_delivery": 0.6784080802039667,
-      "loss_biodist": 0.2691292708570307,
-      "loss_toxic": 0.08434606292708353
+      "loss": 2.418043158271096,
+      "loss_size": 0.25799565559083765,
+      "loss_pdi": 0.5286295684901151,
+      "loss_ee": 0.73835120417855,
+      "loss_delivery": 0.5078353543173183,
+      "loss_biodist": 0.2939920425415039,
+      "loss_toxic": 0.09123929352922873
     },
     {
-      "loss": 2.481573982672258,
-      "loss_size": 0.2506237829273397,
-      "loss_pdi": 0.5553325793959878,
-      "loss_ee": 0.702888163653287,
-      "loss_delivery": 0.6528384658423337,
-      "loss_biodist": 0.25332851166074927,
-      "loss_toxic": 0.06656241179867224
+      "loss": 2.294130650433627,
+      "loss_size": 0.20914554325017062,
+      "loss_pdi": 0.5159178945151243,
+      "loss_ee": 0.7331724437800321,
+      "loss_delivery": 0.50414734875614,
+      "loss_biodist": 0.2559647980061444,
+      "loss_toxic": 0.07578264227644964
     },
     {
-      "loss": 2.328354239463806,
-      "loss_size": 0.20882186158136887,
-      "loss_pdi": 0.5448922569101508,
-      "loss_ee": 0.7051796858960931,
-      "loss_delivery": 0.5663092088970271,
-      "loss_biodist": 0.23630927367643875,
-      "loss_toxic": 0.06684201947328719
+      "loss": 2.260723189874129,
+      "loss_size": 0.21194299920038742,
+      "loss_pdi": 0.51129734787074,
+      "loss_ee": 0.71018939668482,
+      "loss_delivery": 0.506003974513574,
+      "loss_biodist": 0.24507361785932022,
+      "loss_toxic": 0.07621594924818385
     },
     {
-      "loss": 2.1841621182181616,
-      "loss_size": 0.20124886184930801,
-      "loss_pdi": 0.5195452462543141,
-      "loss_ee": 0.6754790571602908,
-      "loss_delivery": 0.5004782859574665,
-      "loss_biodist": 0.2269973118196834,
-      "loss_toxic": 0.06041339209133929
+      "loss": 2.180013732476668,
+      "loss_size": 0.21782933243296362,
+      "loss_pdi": 0.5094848789952018,
+      "loss_ee": 0.6991291533816945,
+      "loss_delivery": 0.4610004154118625,
+      "loss_biodist": 0.22516351396387274,
+      "loss_toxic": 0.06740644057704644
     },
     {
-      "loss": 2.218748081814159,
-      "loss_size": 0.1988528309897943,
-      "loss_pdi": 0.5231576589020815,
-      "loss_ee": 0.6777869246222756,
-      "loss_delivery": 0.5498508228497072,
-      "loss_biodist": 0.20798503336581317,
-      "loss_toxic": 0.06111483716151931
+      "loss": 2.131091995672746,
+      "loss_size": 0.22081551971760663,
+      "loss_pdi": 0.4984923790801655,
+      "loss_ee": 0.6744041009382769,
+      "loss_delivery": 0.44286114688624034,
+      "loss_biodist": 0.2276345125653527,
+      "loss_toxic": 0.0668843225999312
     },
     {
-      "loss": 2.1877676248550415,
-      "loss_size": 0.20004721256819638,
-      "loss_pdi": 0.5174875286492434,
-      "loss_ee": 0.6804409514773976,
-      "loss_delivery": 0.5348355831070379,
-      "loss_biodist": 0.2000916667959907,
-      "loss_toxic": 0.054864687167785385
+      "loss": 2.075855114243247,
+      "loss_size": 0.1967273937030272,
+      "loss_pdi": 0.4761518023230813,
+      "loss_ee": 0.6580501876094125,
+      "loss_delivery": 0.4651151258837093,
+      "loss_biodist": 0.22098628905686465,
+      "loss_toxic": 0.05882429365407337
     },
     {
-      "loss": 2.0138474811207163,
-      "loss_size": 0.19752372259443457,
-      "loss_pdi": 0.4868703836744482,
-      "loss_ee": 0.6450781768018549,
-      "loss_delivery": 0.44826274839314545,
-      "loss_biodist": 0.18033211339603772,
-      "loss_toxic": 0.05578036225316199
+      "loss": 2.070832209153609,
+      "loss_size": 0.22619221427223898,
+      "loss_pdi": 0.46735330332409253,
+      "loss_ee": 0.658088050105355,
+      "loss_delivery": 0.46364551173015073,
+      "loss_biodist": 0.1992648494514552,
+      "loss_toxic": 0.0562882690097798
     },
     {
-      "loss": 2.011049357327548,
-      "loss_size": 0.18426605652679096,
-      "loss_pdi": 0.4944283068180084,
-      "loss_ee": 0.6483827070756392,
-      "loss_delivery": 0.4382695244117217,
-      "loss_biodist": 0.19008470394394614,
-      "loss_toxic": 0.05561802129853855
+      "loss": 2.0497749502008613,
+      "loss_size": 0.2018005665053021,
+      "loss_pdi": 0.44933846592903137,
+      "loss_ee": 0.6409396637569774,
+      "loss_delivery": 0.4817944710904902,
+      "loss_biodist": 0.21351950141516599,
+      "loss_toxic": 0.0623823038556359
     },
     {
-      "loss": 2.0496722134676846,
-      "loss_size": 0.18194164742122998,
-      "loss_pdi": 0.4917480132796548,
-      "loss_ee": 0.6405326967889612,
-      "loss_delivery": 0.5158521959727461,
-      "loss_biodist": 0.1772592243823138,
-      "loss_toxic": 0.042338407937098636
+      "loss": 1.998817953196439,
+      "loss_size": 0.19020642136985605,
+      "loss_pdi": 0.4579390991817821,
+      "loss_ee": 0.6345709101720289,
+      "loss_delivery": 0.45773128284649417,
+      "loss_biodist": 0.20633393864740024,
+      "loss_toxic": 0.052036324177276
     },
     {
-      "loss": 1.9639968113465742,
-      "loss_size": 0.1595074331218546,
-      "loss_pdi": 0.48567668145353143,
-      "loss_ee": 0.630230272358114,
-      "loss_delivery": 0.4618720757690343,
-      "loss_biodist": 0.17742999304424634,
-      "loss_toxic": 0.0492804107171568
-    },
-    {
-      "loss": 1.9399811571294612,
-      "loss_size": 0.18978103656660428,
-      "loss_pdi": 0.48150675947015936,
-      "loss_ee": 0.6243780959736217,
-      "loss_delivery": 0.4384633018211885,
-      "loss_biodist": 0.16038654270497235,
-      "loss_toxic": 0.045465379022061825
+      "loss": 1.9732873006300493,
+      "loss_size": 0.18214904246005145,
+      "loss_pdi": 0.46314583312381397,
+      "loss_ee": 0.6480948545716025,
+      "loss_delivery": 0.43798652359030465,
+      "loss_biodist": 0.19162344119765543,
+      "loss_toxic": 0.05028762956234542
     }
   ],
   "val": [
     {
-      "loss": 8.991046905517578,
-      "loss_size": 4.490478515625,
-      "loss_pdi": 0.9734575947125753,
-      "loss_ee": 0.9441032012303671,
-      "loss_delivery": 1.2459152142206829,
-      "loss_biodist": 0.9957353870073954,
-      "loss_toxic": 0.34135735034942627,
+      "loss": 11.670351346333822,
+      "loss_size": 7.447449843088786,
+      "loss_pdi": 1.082938313484192,
+      "loss_ee": 0.9422469735145569,
+      "loss_delivery": 0.7185012102127075,
+      "loss_biodist": 0.945093979438146,
+      "loss_toxic": 0.5341211756070455,
       "acc_pdi": 0.8484848484848485,
       "acc_ee": 0.8333333333333334,
       "acc_toxic": 1.0
     },
     {
-      "loss": 3.4952890078226724,
-      "loss_size": 0.27914075056711835,
-      "loss_pdi": 0.6866898337999979,
-      "loss_ee": 0.7263152599334717,
-      "loss_delivery": 0.7405761281649271,
-      "loss_biodist": 0.9172844588756561,
-      "loss_toxic": 0.14528251190980276,
+      "loss": 3.7281323273976645,
+      "loss_size": 0.4031377931435903,
+      "loss_pdi": 0.7136062383651733,
+      "loss_ee": 0.7757165829340616,
+      "loss_delivery": 0.6889261901378632,
+      "loss_biodist": 0.8803721169630686,
+      "loss_toxic": 0.2663734555244446,
       "acc_pdi": 0.8484848484848485,
       "acc_ee": 0.8333333333333334,
       "acc_toxic": 1.0
     },
     {
-      "loss": 2.8096923430760703,
-      "loss_size": 0.04479753350218137,
-      "loss_pdi": 0.48064420620600384,
-      "loss_ee": 0.6685001452763876,
-      "loss_delivery": 0.7262534300486246,
-      "loss_biodist": 0.827881266673406,
-      "loss_toxic": 0.06161577875415484,
+      "loss": 2.996154228846232,
+      "loss_size": 0.08968868106603622,
+      "loss_pdi": 0.5251734455426534,
+      "loss_ee": 0.705430785814921,
+      "loss_delivery": 0.7259814739227295,
+      "loss_biodist": 0.8580058316389719,
+      "loss_toxic": 0.09187404563029607,
       "acc_pdi": 0.8484848484848485,
       "acc_ee": 0.8333333333333334,
       "acc_toxic": 1.0
     },
     {
-      "loss": 2.756531000137329,
-      "loss_size": 0.07622841248909633,
-      "loss_pdi": 0.463245431582133,
-      "loss_ee": 0.669317622979482,
-      "loss_delivery": 0.6688057581583658,
-      "loss_biodist": 0.8275523781776428,
-      "loss_toxic": 0.05138145387172699,
+      "loss": 3.0241867701212564,
+      "loss_size": 0.06373066206773122,
+      "loss_pdi": 0.5075281461079916,
+      "loss_ee": 0.6890556613604227,
+      "loss_delivery": 0.8740459084510803,
+      "loss_biodist": 0.8290574749310812,
+      "loss_toxic": 0.060768917202949524,
       "acc_pdi": 0.8484848484848485,
       "acc_ee": 0.8333333333333334,
       "acc_toxic": 1.0
     },
     {
-      "loss": 2.9385202328364053,
-      "loss_size": 0.1038126324613889,
-      "loss_pdi": 0.46456684668858844,
-      "loss_ee": 0.705346941947937,
-      "loss_delivery": 0.7891722718874613,
-      "loss_biodist": 0.82027334968249,
-      "loss_toxic": 0.055348185201485954,
+      "loss": 3.2901877562204995,
+      "loss_size": 0.20215384662151337,
+      "loss_pdi": 0.5433482428391775,
+      "loss_ee": 0.7523069183031718,
+      "loss_delivery": 0.9290379285812378,
+      "loss_biodist": 0.8056556979815165,
+      "loss_toxic": 0.057684975365797676,
       "acc_pdi": 0.8484848484848485,
       "acc_ee": 0.8333333333333334,
       "acc_toxic": 1.0
     },
     {
-      "loss": 2.8632373809814453,
-      "loss_size": 0.056236049781243004,
-      "loss_pdi": 0.44324543078740436,
-      "loss_ee": 0.6018350621064504,
-      "loss_delivery": 0.9420756896336874,
-      "loss_biodist": 0.755872001250585,
-      "loss_toxic": 0.06397318094968796,
+      "loss": 3.11362091700236,
+      "loss_size": 0.11133595556020737,
+      "loss_pdi": 0.5679469505945841,
+      "loss_ee": 0.8252793351809183,
+      "loss_delivery": 0.8343843619028727,
+      "loss_biodist": 0.7218613227208456,
+      "loss_toxic": 0.052812947581211724,
       "acc_pdi": 0.8484848484848485,
-      "acc_ee": 0.8333333333333334,
+      "acc_ee": 0.8181818181818182,
       "acc_toxic": 1.0
     },
     {
-      "loss": 2.993934949239095,
-      "loss_size": 0.08202650646368663,
-      "loss_pdi": 0.508717954158783,
-      "loss_ee": 0.7680216828982035,
-      "loss_delivery": 0.8612345655759176,
-      "loss_biodist": 0.7118468681971232,
-      "loss_toxic": 0.0620873523876071,
-      "acc_pdi": 0.8484848484848485,
-      "acc_ee": 0.48484848484848486,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 2.8506224155426025,
-      "loss_size": 0.03146359619374076,
-      "loss_pdi": 0.5297882954279581,
-      "loss_ee": 0.7625004649162292,
-      "loss_delivery": 0.802927960952123,
-      "loss_biodist": 0.6721820036570231,
-      "loss_toxic": 0.05176017774889866,
-      "acc_pdi": 0.8333333333333334,
-      "acc_ee": 0.4696969696969697,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 2.8424178759256997,
-      "loss_size": 0.08800932268301646,
-      "loss_pdi": 0.540622721115748,
-      "loss_ee": 0.839806874593099,
-      "loss_delivery": 0.6362804472446442,
-      "loss_biodist": 0.6938393215338389,
-      "loss_toxic": 0.043859192014982305,
+      "loss": 3.247321446736654,
+      "loss_size": 0.09277657171090443,
+      "loss_pdi": 0.6567125717798868,
+      "loss_ee": 1.0444208979606628,
+      "loss_delivery": 0.7062844236691793,
+      "loss_biodist": 0.6986102362473806,
+      "loss_toxic": 0.04851680745681127,
       "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.168424208958944,
+      "loss_size": 0.05177713930606842,
+      "loss_pdi": 0.5932339330514272,
+      "loss_ee": 0.968136191368103,
+      "loss_delivery": 0.7594618300596873,
+      "loss_biodist": 0.7631273567676544,
+      "loss_toxic": 0.032687741021315254,
+      "acc_pdi": 0.8484848484848485,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.1226733525594077,
+      "loss_size": 0.16442706187566122,
+      "loss_pdi": 0.4861932198206584,
+      "loss_ee": 0.8927785356839498,
+      "loss_delivery": 0.809323231379191,
+      "loss_biodist": 0.7391951779524485,
+      "loss_toxic": 0.030756143853068352,
+      "acc_pdi": 0.8484848484848485,
       "acc_ee": 0.4696969696969697,
       "acc_toxic": 1.0
     },
     {
-      "loss": 2.9927186171213784,
-      "loss_size": 0.04866213048808277,
-      "loss_pdi": 0.5090581774711609,
-      "loss_ee": 0.8101900815963745,
-      "loss_delivery": 0.8959137797355652,
-      "loss_biodist": 0.697561984260877,
-      "loss_toxic": 0.031332316963622965,
+      "loss": 3.4750285943349204,
+      "loss_size": 0.07188746457298596,
+      "loss_pdi": 0.635328451792399,
+      "loss_ee": 1.0510863463083904,
+      "loss_delivery": 0.9019280473391215,
+      "loss_biodist": 0.7850770453612009,
+      "loss_toxic": 0.02972123461465041,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.3787878787878788,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.3856334686279297,
+      "loss_size": 0.0857744167248408,
+      "loss_pdi": 0.5498887598514557,
+      "loss_ee": 0.9292206565539042,
+      "loss_delivery": 1.003569980462392,
+      "loss_biodist": 0.7962689697742462,
+      "loss_toxic": 0.020910644593338173,
       "acc_pdi": 0.8333333333333334,
       "acc_ee": 0.3787878787878788,
       "acc_toxic": 1.0
     },
     {
-      "loss": 2.9314780632654824,
-      "loss_size": 0.06920161470770836,
-      "loss_pdi": 0.5283850431442261,
-      "loss_ee": 0.8220112522443136,
-      "loss_delivery": 0.7353424032529196,
-      "loss_biodist": 0.7471547623475393,
-      "loss_toxic": 0.029382963587219518,
-      "acc_pdi": 0.8181818181818182,
-      "acc_ee": 0.3787878787878788,
+      "loss": 3.594546397527059,
+      "loss_size": 0.058230139315128326,
+      "loss_pdi": 0.7094775040944418,
+      "loss_ee": 1.134681224822998,
+      "loss_delivery": 0.8755488395690918,
+      "loss_biodist": 0.7928893665472666,
+      "loss_toxic": 0.023719362293680508,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.21212121212121213,
       "acc_toxic": 1.0
     },
     {
-      "loss": 2.9559961954752603,
-      "loss_size": 0.05500282160937786,
-      "loss_pdi": 0.516402949889501,
-      "loss_ee": 0.8015920122464498,
-      "loss_delivery": 0.7959124445915222,
-      "loss_biodist": 0.7596821735302607,
-      "loss_toxic": 0.0274037744384259,
-      "acc_pdi": 0.8181818181818182,
-      "acc_ee": 0.3787878787878788,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 2.9268057346343994,
-      "loss_size": 0.07737081746260326,
-      "loss_pdi": 0.5241717199484507,
-      "loss_ee": 0.835136612256368,
-      "loss_delivery": 0.7045041720072428,
-      "loss_biodist": 0.7396295169989268,
-      "loss_toxic": 0.0459929151305308,
-      "acc_pdi": 0.8181818181818182,
-      "acc_ee": 0.3787878787878788,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 2.9870657920837402,
-      "loss_size": 0.08099805439511935,
-      "loss_pdi": 0.5164442459742228,
-      "loss_ee": 0.8418577512105306,
-      "loss_delivery": 0.7609673937161764,
-      "loss_biodist": 0.7428288658459982,
-      "loss_toxic": 0.04396952743021151,
-      "acc_pdi": 0.8181818181818182,
-      "acc_ee": 0.3787878787878788,
-      "acc_toxic": 1.0
-    },
-    {
-      "loss": 2.979916493097941,
-      "loss_size": 0.09317472080389659,
-      "loss_pdi": 0.514099915822347,
-      "loss_ee": 0.8726487557093302,
-      "loss_delivery": 0.6872047583262125,
-      "loss_biodist": 0.7711265037457148,
-      "loss_toxic": 0.04166170318300525,
-      "acc_pdi": 0.8181818181818182,
+      "loss": 3.34331480662028,
+      "loss_size": 0.09952588627735774,
+      "loss_pdi": 0.5444782872994741,
+      "loss_ee": 0.9434934655825297,
+      "loss_delivery": 0.9477877616882324,
+      "loss_biodist": 0.7897172272205353,
+      "loss_toxic": 0.018312191901107628,
+      "acc_pdi": 0.7878787878787878,
       "acc_ee": 0.36363636363636365,
       "acc_toxic": 1.0
     },
     {
-      "loss": 2.937371532122294,
-      "loss_size": 0.06330848174790542,
-      "loss_pdi": 0.5085045297940572,
-      "loss_ee": 0.8129515051841736,
-      "loss_delivery": 0.7455949584643046,
-      "loss_biodist": 0.7728735754887263,
-      "loss_toxic": 0.03413854034927984,
+      "loss": 3.4212222894032798,
+      "loss_size": 0.08121616393327713,
+      "loss_pdi": 0.5517565310001373,
+      "loss_ee": 1.0685155193010967,
+      "loss_delivery": 0.874200721581777,
+      "loss_biodist": 0.828464408715566,
+      "loss_toxic": 0.017068898615737755,
       "acc_pdi": 0.8181818181818182,
-      "acc_ee": 0.3787878787878788,
+      "acc_ee": 0.25757575757575757,
       "acc_toxic": 1.0
     },
     {
-      "loss": 3.0136941274007163,
-      "loss_size": 0.0849883034825325,
-      "loss_pdi": 0.51878755291303,
-      "loss_ee": 0.785295327504476,
-      "loss_delivery": 0.8036739627520243,
-      "loss_biodist": 0.7789742400248846,
-      "loss_toxic": 0.041974871419370174,
-      "acc_pdi": 0.8181818181818182,
-      "acc_ee": 0.3787878787878788,
+      "loss": 3.6395487785339355,
+      "loss_size": 0.07888962080081303,
+      "loss_pdi": 0.5913220842679342,
+      "loss_ee": 1.0437468489011128,
+      "loss_delivery": 1.0660852392514546,
+      "loss_biodist": 0.8421931266784668,
+      "loss_toxic": 0.01731194742023945,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.2727272727272727,
       "acc_toxic": 1.0
     },
     {
-      "loss": 2.9988585313161216,
-      "loss_size": 0.053379556785027184,
-      "loss_pdi": 0.5192934771378835,
-      "loss_ee": 0.838801383972168,
-      "loss_delivery": 0.7629345854123434,
-      "loss_biodist": 0.7728350758552551,
-      "loss_toxic": 0.05161439681736132,
-      "acc_pdi": 0.8181818181818182,
-      "acc_ee": 0.36363636363636365,
+      "loss": 3.5140305360158286,
+      "loss_size": 0.0599971575041612,
+      "loss_pdi": 0.5561938285827637,
+      "loss_ee": 1.0674984057744343,
+      "loss_delivery": 0.9739653070767721,
+      "loss_biodist": 0.8400343159834543,
+      "loss_toxic": 0.01634151643762986,
+      "acc_pdi": 0.7424242424242424,
+      "acc_ee": 0.22727272727272727,
       "acc_toxic": 1.0
     },
     {
-      "loss": 2.9807325998942056,
-      "loss_size": 0.055528260146578155,
-      "loss_pdi": 0.5009754200776418,
-      "loss_ee": 0.8306655287742615,
-      "loss_delivery": 0.7751240134239197,
-      "loss_biodist": 0.7733635157346725,
-      "loss_toxic": 0.045075801705631115,
-      "acc_pdi": 0.8181818181818182,
-      "acc_ee": 0.36363636363636365,
+      "loss": 3.628753344217936,
+      "loss_size": 0.08887146785855293,
+      "loss_pdi": 0.582503984371821,
+      "loss_ee": 1.114095131556193,
+      "loss_delivery": 0.9745903412501017,
+      "loss_biodist": 0.8516232868035635,
+      "loss_toxic": 0.017069284183283646,
+      "acc_pdi": 0.7121212121212122,
+      "acc_ee": 0.19696969696969696,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.6391177972157798,
+      "loss_size": 0.08282352735598882,
+      "loss_pdi": 0.581497848033905,
+      "loss_ee": 1.141685386498769,
+      "loss_delivery": 0.9548555612564087,
+      "loss_biodist": 0.8643284440040588,
+      "loss_toxic": 0.013927079737186432,
+      "acc_pdi": 0.7424242424242424,
+      "acc_ee": 0.16666666666666666,
       "acc_toxic": 1.0
     }
   ]
diff --git a/models/finetune_cv/fold_4/model.pt b/models/finetune_cv/fold_4/model.pt
index a8c8de6..c4f023c 100644
Binary files a/models/finetune_cv/fold_4/model.pt and b/models/finetune_cv/fold_4/model.pt differ
diff --git a/models/finetune_cv/test_results.json b/models/finetune_cv/test_results.json
index 72d2e3f..0014c1a 100644
--- a/models/finetune_cv/test_results.json
+++ b/models/finetune_cv/test_results.json
@@ -5,15 +5,15 @@
       "n_samples": 95,
       "size": {
         "n": 95,
-        "rmse": 0.632240295394253,
-        "mae": 0.4674149764211554,
-        "r2": -0.13819694158244777
+        "rmse": 0.5909209144067168,
+        "mae": 0.376253614927593,
+        "r2": 0.005712927161997228
       },
       "delivery": {
         "n": 66,
-        "rmse": 1.3553486689823926,
-        "mae": 0.48114709207562334,
-        "r2": -0.008234200686471516
+        "rmse": 1.3280577883458438,
+        "mae": 0.5195405159964028,
+        "r2": 0.03195999739694366
       },
       "pdi": {
         "n": 95,
@@ -24,10 +24,10 @@
       },
       "ee": {
         "n": 95,
-        "accuracy": 0.6631578947368421,
-        "precision": 0.22580645161290322,
-        "recall": 0.328125,
-        "f1": 0.267515923566879
+        "accuracy": 0.6736842105263158,
+        "precision": 0.22456140350877193,
+        "recall": 0.3333333333333333,
+        "f1": 0.26834381551362685
       },
       "toxic": {
         "n": 66,
@@ -38,8 +38,8 @@
       },
       "biodist": {
         "n": 66,
-        "kl_divergence": 0.8735123101690443,
-        "js_divergence": 0.2203766048579219
+        "kl_divergence": 0.851655784204727,
+        "js_divergence": 0.21404831573756974
       }
     },
     {
@@ -47,15 +47,15 @@
       "n_samples": 195,
       "size": {
         "n": 193,
-        "rmse": 0.42622752144538556,
-        "mae": 0.24566329575573226,
-        "r2": 0.0482086242002292
+        "rmse": 0.4425801645813746,
+        "mae": 0.26432527161632796,
+        "r2": -0.026225211870033682
       },
       "delivery": {
         "n": 123,
-        "rmse": 0.742899240869997,
-        "mae": 0.5315999669170507,
-        "r2": -0.03140039086191759
+        "rmse": 0.7771322048436382,
+        "mae": 0.6133777339870822,
+        "r2": -0.128644776760948
       },
       "pdi": {
         "n": 195,
@@ -80,8 +80,8 @@
       },
       "biodist": {
         "n": 123,
-        "kl_divergence": 1.218490162626268,
-        "js_divergence": 0.3069412988090679
+        "kl_divergence": 0.9336461102028436,
+        "js_divergence": 0.24870266224462317
       }
     },
     {
@@ -89,15 +89,15 @@
       "n_samples": 51,
       "size": {
         "n": 51,
-        "rmse": 0.241909571406037,
-        "mae": 0.20043573192521638,
-        "r2": -0.43487628292073
+        "rmse": 0.6473513298834871,
+        "mae": 0.5600235602434944,
+        "r2": -9.27515642706235
       },
       "delivery": {
         "n": 44,
-        "rmse": 0.7564153649581582,
-        "mae": 0.6047130756302398,
-        "r2": -0.4226486727361405
+        "rmse": 0.7721077356414991,
+        "mae": 0.6167582499593581,
+        "r2": -0.4822886602727561
       },
       "pdi": {
         "n": 51,
@@ -122,8 +122,8 @@
       },
       "biodist": {
         "n": 45,
-        "kl_divergence": 1.0902737810034395,
-        "js_divergence": 0.27372590260311874
+        "kl_divergence": 1.1049896129018548,
+        "js_divergence": 0.25485248115851133
       }
     },
     {
@@ -131,15 +131,15 @@
       "n_samples": 66,
       "size": {
         "n": 66,
-        "rmse": 0.2857872773679936,
-        "mae": 0.22075237649859805,
-        "r2": -0.5674047032859011
+        "rmse": 0.2407212117920812,
+        "mae": 0.19363613562150436,
+        "r2": -0.11204941379936861
       },
       "delivery": {
         "n": 62,
-        "rmse": 1.0291312965402932,
-        "mae": 0.7422042032328224,
-        "r2": -0.7148264932933832
+        "rmse": 1.0041711455927012,
+        "mae": 0.7132550483914993,
+        "r2": -0.63265374674746
       },
       "pdi": {
         "n": 66,
@@ -150,10 +150,10 @@
       },
       "ee": {
         "n": 66,
-        "accuracy": 0.18181818181818182,
-        "precision": 0.203921568627451,
-        "recall": 0.2707070707070707,
-        "f1": 0.12297410192147035
+        "accuracy": 0.8181818181818182,
+        "precision": 0.27692307692307694,
+        "recall": 0.32727272727272727,
+        "f1": 0.3
       },
       "toxic": {
         "n": 62,
@@ -164,8 +164,8 @@
       },
       "biodist": {
         "n": 62,
-        "kl_divergence": 0.9434472801013084,
-        "js_divergence": 0.20391570021898991
+        "kl_divergence": 0.9677978984139058,
+        "js_divergence": 0.2020309307244639
       }
     },
     {
@@ -173,15 +173,15 @@
       "n_samples": 27,
       "size": {
         "n": 27,
-        "rmse": 0.2271495001169846,
-        "mae": 0.18753767013549805,
-        "r2": -0.19441156195074893
+        "rmse": 0.23392834445509142,
+        "mae": 0.19066280788845485,
+        "r2": -0.2667651950955112
       },
       "delivery": {
         "n": 15,
-        "rmse": 1.993006453768918,
-        "mae": 1.3779302000999452,
-        "r2": -0.3411461507368889
+        "rmse": 1.9603892288630869,
+        "mae": 1.3892907698949177,
+        "r2": -0.29760739742916287
       },
       "pdi": {
         "n": 27,
@@ -206,23 +206,23 @@
       },
       "biodist": {
         "n": 15,
-        "kl_divergence": 0.7615404990024607,
-        "js_divergence": 0.20916908426734354
+        "kl_divergence": 0.9389607012315264,
+        "js_divergence": 0.2470218476598176
       }
     }
   ],
   "summary_stats": {
     "size": {
-      "rmse_mean": 0.36266283314613074,
-      "rmse_std": 0.15203127472757474,
-      "r2_mean": -0.2573361731079197,
-      "r2_std": 0.2187118059634264
+      "rmse_mean": 0.43110039302375025,
+      "rmse_std": 0.17179051271013462,
+      "r2_mean": -1.9348966641330534,
+      "r2_std": 3.6713441784129
     },
     "delivery": {
-      "rmse_mean": 1.1753602050239518,
-      "rmse_std": 0.46580283242073095,
-      "r2_mean": -0.30365118166296035,
-      "r2_std": 0.2630677092396549
+      "rmse_mean": 1.1683716206573538,
+      "rmse_std": 0.4449374578352648,
+      "r2_mean": -0.30184691676267666,
+      "r2_std": 0.23809090378746706
     },
     "pdi": {
       "accuracy_mean": 0.7875890604063979,
@@ -231,10 +231,10 @@
       "f1_std": 0.08529090446864619
     },
     "ee": {
-      "accuracy_mean": 0.5402437489124795,
-      "accuracy_std": 0.22467627690136344,
-      "f1_mean": 0.2281738334487203,
-      "f1_std": 0.063019179670124
+      "accuracy_mean": 0.6696217393431015,
+      "accuracy_std": 0.15503740047242787,
+      "f1_mean": 0.2637445914537758,
+      "f1_std": 0.039213602228007696
     },
     "toxic": {
       "accuracy_mean": 0.9490006447453256,
@@ -243,26 +243,26 @@
       "f1_std": 0.26169039387919035
     },
     "biodist": {
-      "kl_mean": 0.9774528065805042,
-      "kl_std": 0.1608761675751487,
-      "js_mean": 0.24282571815128842,
-      "js_std": 0.04053726747029075
+      "kl_mean": 0.9594100213909715,
+      "kl_std": 0.08240959093662605,
+      "js_mean": 0.23333124750499712,
+      "js_std": 0.021158533549255752
     }
   },
   "overall": {
     "size": {
       "n_samples": 432,
-      "mse": 0.19167728610863985,
-      "rmse": 0.43780964597486866,
-      "mae": 0.2816500812768936,
-      "r2": -0.04410163027802061
+      "mse": 0.22604480336185886,
+      "rmse": 0.47544169291497657,
+      "mae": 0.3084443360567093,
+      "r2": -0.2313078534105617
     },
     "delivery": {
       "n_samples": 310,
-      "mse": 1.095306046771274,
-      "rmse": 1.0465687014101244,
-      "mae": 0.6143080417337197,
-      "r2": -0.1024184074306409
+      "mse": 1.0873755440675295,
+      "rmse": 1.0427730069710903,
+      "mae": 0.6513989447841361,
+      "r2": -0.09443640807387799
     },
     "pdi": {
       "n_samples": 434,
@@ -273,10 +273,10 @@
     },
     "ee": {
       "n_samples": 434,
-      "accuracy": 0.4976958525345622,
-      "precision": 0.21063404810247777,
-      "recall": 0.2839160839160839,
-      "f1": 0.23684873775319115
+      "accuracy": 0.5967741935483871,
+      "precision": 0.1993841416474211,
+      "recall": 0.33205128205128204,
+      "f1": 0.24915824915824913
     },
     "toxic": {
       "n_samples": 313,
@@ -287,8 +287,8 @@
     },
     "biodist": {
       "n_samples": 311,
-      "kl_divergence": 1.0498561462079121,
-      "js_divergence": 0.25851000311532496
+      "kl_divergence": 0.9481034280166569,
+      "js_divergence": 0.23285280825310384
     }
   }
 }
\ No newline at end of file