fix the bug for E only mode for save logic

MDIL-SNU · Aug 24, 2018 · 197ce50 · 197ce50
1 parent 745840e
commit 197ce50
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/simple_nn/models/neural_network.py b/simple_nn/models/neural_network.py
@@ -619,7 +619,8 @@ def train(self, user_optimizer=None, user_atomic_weights_function=None):
 
                     # Temp saving
                     #if (epoch+1) % self.inputs['save_interval'] == 0:
-                    if save_stack > self.inputs['save_interval'] and prev_eloss > eloss and prev_floss > floss:
+                    if save_stack > self.inputs['save_interval'] and prev_eloss > eloss and \
+                       ((prev_floss > floss) or floss == 0.):
                         self._save(sess, saver)
                         prev_eloss = eloss
                         prev_floss = floss