Skip to content

Commit

Permalink
🎨 Format Python code with psf/black
Browse files Browse the repository at this point in the history
  • Loading branch information
mzouink committed Mar 17, 2024
1 parent 636c416 commit bf1acfa
Showing 1 changed file with 4 additions and 5 deletions.
9 changes: 4 additions & 5 deletions src/cellmap_models/pytorch/cosem/load_checkpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,7 @@
from cellmap_models import download_url_to_file



def download_checkpoint(checkpoint_name: str,local_folder: Path):
def download_checkpoint(checkpoint_name: str, local_folder: Path):
"""
download models checkpoint from s3 bucket.
Expand All @@ -13,16 +12,16 @@ def download_checkpoint(checkpoint_name: str,local_folder: Path):
return:
checkpoint_path (Path): Path to the downloaded checkpoint.
"""
from . import models_dict, models_list, model_names # avoid circular import
from . import models_dict, models_list # avoid circular import

# Make sure the checkpoint exists
if checkpoint_name not in models_list:
raise ValueError(
f"Checkpoint {checkpoint_name} not found. Available checkpoints: {models_list}"
)

checkpoint_path = Path(
local_folder/ Path(checkpoint_name.replace(".", "_"))
local_folder / Path(checkpoint_name.replace(".", "_"))
).with_suffix(".pth")
if not checkpoint_path.exists():
url = models_dict[checkpoint_name]
Expand Down

0 comments on commit bf1acfa

Please sign in to comment.