Skip to content

Commit

Permalink
Apply suggestions from code review
Browse files Browse the repository at this point in the history
Co-authored-by: Justus Schock <12886177+justusschock@users.noreply.github.com>
  • Loading branch information
Borda and justusschock authored Sep 25, 2020
1 parent 551a980 commit aa52cc4
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions pytorch_lightning/callbacks/model_checkpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,8 @@
import re
import yaml
from copy import deepcopy
from typing import Any, Dict, Optional
from typing import Any, Dict, Optional, Union
from pathlib import Path

import numpy as np
import torch
Expand Down Expand Up @@ -468,7 +469,7 @@ def on_load_checkpoint(self, checkpointed_state: Dict[str, Any]):
self.best_model_score = checkpointed_state["best_model_score"]
self.best_model_path = checkpointed_state["best_model_path"]

def to_yaml(self, save_path=None):
def to_yaml(self, save_path: Optional[Union[str, Path]]=None):
""" Saves `{'checkpoint_name': score}` dict as a YAML file."""
best_k = {k: v.item() for k, v in self.best_k_models.items()}
if save_path is None:
Expand Down

0 comments on commit aa52cc4

Please sign in to comment.