diff --git a/tuning.py b/tuning.py
index 8c2d5e46..61cfd854 100644
--- a/tuning.py
+++ b/tuning.py
@@ -3,10 +3,12 @@
 import torch.cuda
 from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning import Trainer
-from train_pl import *
 from ray import tune
 from ray.tune.integration.pytorch_lightning import TuneReportCallback
 
+from train_pl import *
+
+
 args = PLOptions().parse()
 
 def train_segmentation(config):
@@ -14,14 +16,10 @@ def train_segmentation(config):
         args.__dict__[k] = v
     
     model = MeshSegmenter(args)
-    callback_tune = TuneReportCallback(metrics='val_iou', on="validation_end")
+    callback_tune = TuneReportCallback(metrics=['val_iou', 'val_f1'], on="validation_end")
     callback_lightning = ModelCheckpoint(monitor='val_iou', mode='max', save_top_k=3,
                                          filename='{epoch:02d}-{val_acc_epoch:.2f}', )
 
-    # callback_tune_f1 = TuneReportCallback(metrics='val_f1', on="validation_end")
-    # callback_lightning_f1 = ModelCheckpoint(monitor='val_f1', mode='max', save_top_k=3,
-    #                                      filename='{epoch:02d}-{val_acc_epoch:.2f}', )
-
     trainer = Trainer.from_argparse_args(args, callbacks=[callback_tune, callback_lightning])
     trainer.fit(model)
     torch.cuda.empty_cache()