Merge: [DLRM/PyT] Add validation loss logging to test mode

nv-kkudrynski · nv-kkudrynski · commit 382524929a42 · 2022-07-07T04:06:18.000-07:00
diff --git a/PyTorch/Recommendation/DLRM/dlrm/scripts/main.py b/PyTorch/Recommendation/DLRM/dlrm/scripts/main.py
@@ -509,13 +509,10 @@ def parallelize(model):
 
     if FLAGS.mode == 'test':
         model = parallelize(model)
-        auc = dist_evaluate(model, data_loader_test)
+        auc, valid_loss = dist_evaluate(model, data_loader_test)
 
-        results = {'auc': auc}
+        results = {'best_auc': auc, 'best_validation_loss': valid_loss}
         dllogger.log(data=results, step=tuple())
-
-        if auc is not None:
-            print(f"Finished testing. Test auc {auc:.4f}")
         return
     elif FLAGS.mode == 'inference_benchmark':
         if world_size > 1:
