Update pytorch lightning training

Vchitect · Sep 3, 2024 · e7f9415 · e7f9415
1 parent 1b74d1e
commit e7f9415
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/train_pl.py b/train_pl.py
@@ -37,7 +37,7 @@ def __init__(self, args, logger: logging.Logger):
         requires_grad(self.ema, False)
 
         # Load pretrained model if specified
-        if args.pretrained and args.resume_from_checkpoint is not None:
+        if args.pretrained:
             # Load old checkpoint, only load EMA
             self._load_pretrained_parameters(args)
         self.logging.info(f"Model Parameters: {sum(p.numel() for p in self.model.parameters()):,}")

diff --git a/train_with_img_pl.py b/train_with_img_pl.py
@@ -37,7 +37,7 @@ def __init__(self, args, logger: logging.Logger):
         requires_grad(self.ema, False)
 
         # Load pretrained model if specified
-        if args.pretrained and args.resume_from_checkpoint is not None:
+        if args.pretrained:
             # Load old checkpoint, only load EMA
             self._load_pretrained_parameters(args)
         self.logging.info(f"Model Parameters: {sum(p.numel() for p in self.model.parameters()):,}")