diff --git a/src/napatrackmater/Trackvector.py b/src/napatrackmater/Trackvector.py
index c6cdeaae..35d700cd 100644
--- a/src/napatrackmater/Trackvector.py
+++ b/src/napatrackmater/Trackvector.py
@@ -1406,6 +1406,11 @@ def train_mitosis_neural_net(
                         if total_val_class2 > 0
                         else 0
                     )
+                    loss_class1 = criterion_class1(outputs_class1, labels_class1)
+                    loss_class2 = criterion_class2(outputs_class2, labels_class2)
+
+                    running_val_loss_class1 += loss_class1.item()
+                    running_val_loss_class2 += loss_class2.item()
                     pbar_val.set_postfix(
                         {"Acc Class1": accuracy_class1, "Acc Class2": accuracy_class2}
                     )
@@ -1494,7 +1499,9 @@ def predict_with_model(saved_model_path, saved_model_json, features_array):
         num_classes_class1=num_classes_class1,
         num_classes_class2=num_classes_class2,
     )
-    model.load_state_dict(torch.load(saved_model_path, map_location=torch.device(device)))
+    model.load_state_dict(
+        torch.load(saved_model_path, map_location=torch.device(device))
+    )
     model.to(device)
     model.eval()
 
@@ -1505,11 +1512,9 @@ def predict_with_model(saved_model_path, saved_model_json, features_array):
         outputs_class1, outputs_class2 = model(features_tensor)
         predicted_probs_class1 = torch.softmax(outputs_class1, dim=1)
         predicted_probs_class2 = torch.softmax(outputs_class2, dim=1)
-        
+
         predicted_class1 = torch.argmax(predicted_probs_class1, dim=1).cpu().numpy()
         predicted_class2 = torch.argmax(predicted_probs_class2, dim=1).cpu().numpy()
-        
-        
 
     return predicted_class1, predicted_class2
 
diff --git a/src/napatrackmater/__init__.py b/src/napatrackmater/__init__.py
index c39234a9..416f1773 100644
--- a/src/napatrackmater/__init__.py
+++ b/src/napatrackmater/__init__.py
@@ -14,8 +14,8 @@
     train_mitosis_classifier,
     train_mitosis_neural_net,
     extract_neural_training_data,
-    predict_with_model
-
+    predict_with_model,
+    plot_metrics_from_npz,
 )
 from .DeepEmbeddedClustering import DeepEmbeddedClustering
 from .CloudAutoEncoder import CloudAutoEncoder
@@ -73,11 +73,12 @@ def load_json(fpath):
     "unsupervised_clustering",
     "DeepEmbeddedClustering",
     "load_json",
-    'create_mitosis_training_data',
-    'train_mitosis_classifier',
-    'train_mitosis_neural_net',
-    'extract_neural_training_data',
-    'predict_with_model'
+    "create_mitosis_training_data",
+    "train_mitosis_classifier",
+    "train_mitosis_neural_net",
+    "extract_neural_training_data",
+    "predict_with_model",
+    "plot_metrics_from_npz",
 )
 
 clear_models_and_aliases(DeepEmbeddedClustering, CloudAutoEncoder)
diff --git a/src/napatrackmater/_version.py b/src/napatrackmater/_version.py
index 6556e003..01e6bcfc 100644
--- a/src/napatrackmater/_version.py
+++ b/src/napatrackmater/_version.py
@@ -1,2 +1,2 @@
-__version__ = version = "4.4.7"
-__version_tuple__ = version_tuple = (4, 4, 7)
+__version__ = version = "4.4.8"
+__version_tuple__ = version_tuple = (4, 4, 8)