small updates

2023-05-22 00:14:18 +02:00
parent 52e461ca47
commit 6b81c0d1a6
5 changed files with 202 additions and 7 deletions
@@ -31,7 +31,7 @@ class CheckpointSaver:
        self.top_model_paths = self.top_model_paths[len(self.top_model_paths)-self.top_n:]


-def load_model(model, checkpoint, optimizer=None, scaler=None, iteration=''):
+def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=None, scaler: torch.nn.Module=None, iteration=''):
    # list of checkpoint files in the folder checkpoints
    checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')]
    checkpoint = torch.load(checkpoints[-1])