mitosis training

Kapoorlabs-CAPED · Nov 25, 2023 · 97c5ad8 · 97c5ad8
1 parent 6062572
commit 97c5ad8
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/napatrackmater/Trackvector.py b/src/napatrackmater/Trackvector.py
@@ -1296,6 +1296,11 @@ def train_mitosis_neural_net(features_array, labels_array_class1, labels_array_c
                 loss_class2.backward()
 
                 optimizer.step()
+
+                outputs_class1, outputs_class2 = model(inputs)
+
+                _, predicted_class1 = torch.max(outputs_class1.data, 1)
+                _, predicted_class2 = torch.max(outputs_class2.data, 1)
 
                 running_loss_class1 += loss_class1.item()
                 running_loss_class2 += loss_class2.item()
@@ -1359,7 +1364,7 @@ def train_mitosis_neural_net(features_array, labels_array_class1, labels_array_c
     torch.save(model.state_dict(), save_path + '_mitosis_track_model.pth')
 
 
-    
+
 def plot_metrics_from_npz(npz_file):
     data = np.load(npz_file)