SaTformer: A Space-Time Transformer for Precipitation Nowcasting

Authors: Levi Harris, Tianlong Chen — The University of North Carolina at Chapel Hill

Usage

import torch
from huggingface_hub import hf_hub_download
from src.model.SaTformer.SaTformer import SaTformer

model = SaTformer(
    dim=512,
    num_frames=4,
    num_classes=64,
    image_size=32,
    patch_size=4,
    channels=11,
    depth=12,
    heads=8,
    dim_head=64,
    attn_dropout=0.1,
    ff_dropout=0.1,
    rotary_emb=False,
    attn="ST^2"
)

weights = hf_hub_download(repo_id="leharris3/satformer", filename="sf-64-cls.pt")
model.load_state_dict(torch.load(weights, weights_only=True), strict=False)
model.eval()

with torch.no_grad():
    x = torch.rand(1, 4, 11, 32, 32)  # (batch, frames, channels, H, W)
    logits = model(x)                  # -> [1, 64]

Citation

@article{harris2025satformer,
  title={A Space-Time Transformer for Precipitation Forecasting},
  author={Harris, Levi and Chen, Tianlong},
  journal={arXiv preprint arXiv:2511.11090},
  year={2025}
}