Merge pull request #119 from zhang-haojie/main

Update pytorch lightning training
Vchitect · Sep 3, 2024 · 586a76d · 586a76d
2 parents 1b74d1e + e7f9415
commit 586a76d
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()):,}")