fix bug of DRIT

HelenMao · HelenMao · commit e386c252f059 · 2019-04-23T14:15:19.000+08:00
diff --git a/DRIT-Mode-Seeking/train.py b/DRIT-Mode-Seeking/train.py
@@ -59,7 +59,7 @@ def main():
       total_it += 1
       if total_it >= max_it:
         saver.write_img(-1, model)
-        saver.write_model(-1, model)
+        saver.write_model(-1, max_it, model)
         break
 
     # decay learning rate
