Fix train.py for regression task

tristandeleu · Apr 26, 2020 · deecc93 · deecc93
1 parent 9da961a
commit deecc93
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/train.py b/train.py
@@ -76,10 +76,10 @@ def main(args):
                                        desc=epoch_desc.format(epoch + 1))
 
         # Save best model
-        if (best_value is None) or (('accuracies_after' in results)
-                and (best_value < results['accuracies_after'])):
-            best_value = results['accuracies_after']
-            save_model = True
+        if 'accuracies_after' in results:
+            if (best_value is None) or (best_value < results['accuracies_after']):
+                best_value = results['accuracies_after']
+                save_model = True
         elif (best_value is None) or (best_value > results['mean_outer_loss']):
             best_value = results['mean_outer_loss']
             save_model = True