Added missed returned loss on baseline in experiment.py

tnoe1 · Jul 22, 2022 · 4f3953e · 4f3953e
1 parent 346344a
commit 4f3953e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/experiment.py b/experiment.py
@@ -283,7 +283,7 @@ def main():
     for i in range(0, epochs):
         # start_time = time.time()
         if args.baseline:
-            train_ce(net, id_train_loader, optim, i, id_label_map, device)
+            loss = train_ce(net, id_train_loader, optim, i, id_label_map, device)
         else:
             if args.loss == "margin" and args.detection_type == "LS":
                 loss = train_lm_ls(net, lm, id_train_loader, ood_train_loader, optim, i, id_label_map, device)
@@ -326,7 +326,7 @@ def main():
         if i % 5 == 0 or metric_combined > metric_combined_running_max:
             if metric_combined > metric_combined_running_max:
                 metric_combined_running_max = metric_combined
-                
+
             if args.baseline:
                 directory = "baseline_{}".format(test_method)
                 torch.save({