Nano-JEPA is a lightweight, efficient implementation of the Joint Embedding Predictive Architecture (JEPA) designed for community-driven research.
Trained on a dual-GPU setup achieving ~49% MFU (Model Flops Utilization).
1import torch
2from model import NanoJEPA
3
4# Initialize architecture (ensure this matches your training config)
5model = NanoJEPA(
6 img_size=224,
7 patch_size=16,
8 embed_dim=768,
9 context_depth=12,
10 predictor_depth=4
11)
12
13# Load weights
14state_dict = torch.load("weights/nano-jepa-200M.pth")
15model.load_state_dict(state_dict)
16model.eval()
