def _load_snapshot()

in pai-python-sdk/training/pytorch_ddp/train_src/train_multinode.py [0:0]

• 6 lines of code
• 1 McCabe index (conditional complexity)

    def _load_snapshot(self, snapshot_path):
        loc = f"cuda:{self.local_rank}"
        snapshot = torch.load(snapshot_path, map_location=loc)
        self.model.load_state_dict(snapshot["MODEL_STATE"])
        self.epochs_run = snapshot["EPOCHS_RUN"]
        print(f"Resuming training from snapshot at Epoch {self.epochs_run}")