bug fix

CASIA-IVA-Lab · Sep 18, 2023 · d941c82 · d941c82
1 parent 72642ff
commit d941c82
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/code/train_mvtec.py b/code/train_mvtec.py
@@ -100,9 +100,8 @@ def main(**args):
             del batch_sft
             current_step += 1
             # torch.cuda.empty_cache()
-            current_step += 1
-            if iter_every_epoch % 1000 == 0:
-                agent.save_model(args['save_path'], 0)
+            # if iter_every_epoch % 1000 == 0:
+            #     agent.save_model(args['save_path'], 0)
         # save at the end of the training
         torch.distributed.barrier()
         agent.save_model(args['save_path'], 0)