Reintroduce Atlas with performance flag (#897)

Author: pzharrington

earth2studio/models/px/atlas.py

@@ -162,6 +162,11 @@ class Atlas(torch.nn.Module, AutoModelMixin, PrognosticMixin):
     performs autoregressive timestepping using a full-resolution physical state
     and an internal low-resolution latent state.
 
+    Note
+    ----
+    For best inference performance, set the environment variable `TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=1`.
+    This is on by default in NGC containers, but other environments may need to set it manually.
+
     Badges
     ------
     region:global class:mrf product:wind product:precip product:temp product:atmos year:2026

…medium_range/06_atlas_inference.disabled …es/02_medium_range/06_atlas_inference.pyexamples/02_medium_range/06_atlas_inference.disabled renamed to examples/02_medium_range/06_atlas_inference.py examples/02_medium_range/06_atlas_inference.disabled renamed to examples/02_medium_range/06_atlas_inference.py

@@ -63,6 +63,9 @@
 import os
 
 os.makedirs("outputs", exist_ok=True)
+# Performance optimization for Atlas model
+# This is on by default in NGC containers, but other environments may need to set it manually.
+os.environ["TORCH_ALLOW_TF32_CUBLAS_OVERRIDE"] = "1"
 from dotenv import load_dotenv
 
 load_dotenv()  # TODO: make common example prep function

tox.ini

@@ -39,6 +39,8 @@ setenv =
     # https://docs.jax.dev/en/latest/gpu_memory_allocation.html
     XLA_PYTHON_CLIENT_ALLOCATOR = platform
     XLA_PYTHON_CLIENT_PREALLOCATE = false
+    # Performance optimization for Atlas model
+    TORCH_ALLOW_TF32_CUBLAS_OVERRIDE = 1
     # Unique coverage file per tox env so parallel CI jobs don't overwrite each other
     COVERAGE_FILE = .coverage.{envname}