Fix config for 0.1 deg model

microsoft · Aug 23, 2024 · d1ca223 · d1ca223
1 parent c39cfb1
commit d1ca223
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 5 deletions.
diff --git a/aurora/model/aurora.py b/aurora/model/aurora.py
@@ -13,7 +13,7 @@
 from aurora.model.lora import LoRAMode
 from aurora.model.swin3d import Swin3DTransformerBackbone
 
-__all__ = ["Aurora", "AuroraSmall"]
+__all__ = ["Aurora", "AuroraSmall", "AuroraHighRes"]
 
 
 class Aurora(torch.nn.Module):
@@ -250,3 +250,9 @@ def load_checkpoint(self, repo: str, name: str, strict: bool = True) -> None:
     num_heads=8,
     use_lora=False,
 )
+
+AuroraHighRes = partial(
+    Aurora,
+    encoder_depths=(6, 8, 8),
+    decoder_depths=(8, 8, 6),
+)
diff --git a/docs/api.rst b/docs/api.rst
@@ -22,3 +22,6 @@ Models
 
 .. autoclass:: aurora.AuroraSmall
     :members:
+
+.. autoclass:: aurora.AuroraHighRes
+    :members:
diff --git a/docs/models.md b/docs/models.md
@@ -127,9 +127,9 @@ Aurora 0.1° Fine-Tuned is a high-resolution version of Aurora.
 ### Usage
 
 ```python
-from aurora import Aurora
+from aurora import AuroraHighRes
 
-model = Aurora()
+model = AuroraHighRes()
 model.load_checkpoint("microsoft/aurora", "aurora-0.1-finetuned.ckpt")
 ```
 
@@ -170,8 +170,8 @@ Therefore, you should use the static variables provided in
 you can turn off LoRA to obtain more realistic predictions at the expensive of slightly higher long-term MSE:
 
 ```python
-from aurora import Aurora
+from aurora import AuroraHighRes
 
-model = Aurora(use_lora=False)  # Disable LoRA for more realistic samples.
+model = AuroraHighRes(use_lora=False)  # Disable LoRA for more realistic samples.
 model.load_checkpoint("microsoft/aurora", "aurora-0.1-finetuned.ckpt", strict=False)
 ```