1import torch
2from huggingface_hub import hf_hub_download
3from src.model.SaTformer.SaTformer import SaTformer
4
5model = SaTformer(
6 dim=512,
7 num_frames=4,
8 num_classes=64,
9 image_size=32,
10 patch_size=4,
11 channels=11,
12 depth=12,
13 heads=8,
14 dim_head=64,
15 attn_dropout=0.1,
16 ff_dropout=0.1,
17 rotary_emb=False,
18 attn="ST^2"
19)
20
21weights = hf_hub_download(repo_id="leharris3/satformer", filename="sf-64-cls.pt")
22model.load_state_dict(torch.load(weights, weights_only=True), strict=False)
23model.eval()
24
25with torch.no_grad():
26 x = torch.rand(1, 4, 11, 32, 32) # (batch, frames, channels, H, W)
27 logits = model(x) # -> [1, 64]
1@article{harris2025satformer,
2 title={A Space-Time Transformer for Precipitation Forecasting},
3 author={Harris, Levi and Chen, Tianlong},
4 journal={arXiv preprint arXiv:2511.11090},
5 year={2025}
6}
