FIX: Generalize warmup