diff --git a/pyproject.toml b/pyproject.toml index 5cdf259..9750bed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,8 +60,8 @@ optional-dependencies.dev = [ urls.Repository = "https://github.com/ecmwf-lab/ai-models-aurora" entry-points."ai_models.model".aurora = "ai_models_aurora.model:model" entry-points."ai_models.model"."aurora-0.1-finetuned" = "ai_models_aurora.model:Aurora0p1FineTuned" -entry-points."ai_models.model"."aurora-2.5-finetuned" = "ai_models_aurora.model:Aurora2p5FineTuned" -entry-points."ai_models.model"."aurora-2.5-pretrained" = "ai_models_aurora.model:Aurora2p5Pretrained" +entry-points."ai_models.model"."aurora-2.5-finetuned" = "ai_models_aurora.model:Aurora0p25FineTuned" +entry-points."ai_models.model"."aurora-2.5-pretrained" = "ai_models_aurora.model:Aurora0p25Pretrained" [tool.setuptools_scm] version_file = "src/ai_models_aurora/_version.py" diff --git a/src/ai_models_aurora/model.py b/src/ai_models_aurora/model.py index 8f47e2e..699561c 100644 --- a/src/ai_models_aurora/model.py +++ b/src/ai_models_aurora/model.py @@ -193,7 +193,7 @@ class Aurora0p25Pretrained(Aurora0p25): # https://microsoft.github.io/aurora/models.html#aurora-0-25-fine-tuned -class Aurora025FineTuned(Aurora0p25): +class Aurora0p25FineTuned(Aurora0p25): use_lora = True checkpoint = "aurora-0.25-finetuned.ckpt" @@ -241,7 +241,7 @@ def model(model_version, **kwargs): models = { "0.25-pretrained": Aurora0p25Pretrained, - "0.25-finetuned": Aurora025FineTuned, + "0.25-finetuned": Aurora0p25FineTuned, "0.1-finetuned": Aurora0p1FineTuned, "default": Aurora0p1FineTuned, "latest": Aurora0p1FineTuned, # Backward compatibility