# I Have 10 Minutes to Train an AI Model. Here’s Exactly What Happened. > Source: > Published: 2026-07-11 07:29:03+00:00 A few days ago I sat down, opened a terminal, and gave myself a stupid little constraint: train a real model, on a real dataset, with real eval numbers, in the time it takes to make and drink a cup of chai. No pre-warmed environment. No “oh I already have the CUDA drivers set up.” Cold start, on a rented GPU I’d never touched before, using [Crunr](https://crunr.com/) (the CLI I’ve been building) to get from pip install to a finished, evaluated model. I recorded the whole thing. If you want the unscripted version, it’s here: [10 Minutes to Train an AI Model](https://youtu.be/xdqmuLJzZJk?si=sbcrXj0-sZwk8XaK). This post is the write-up: the code, the numbers, and the two or three things that actually matter if you want to repeat this yourself. Nothing exotic. I fine-tuned distilbert-base-uncased on AG News, a 4-class topic classification dataset (World, Sports, Business, Sci/Tech). It's a well-worn benchmark on purpose. When you're testing infrastructure, the last thing you want is a dataset that's also fighting you. The script does five things: That baseline step matters more than people give it credit for. It’s easy to fine-tune a model, see 94% accuracy, and feel great about it. It’s a different thing to see the untrained model sitting at 23.5% first and watch the actual delta your training run bought you. ``` torch>=2.0.0transformers>=4.38.0datasets>=2.18.0scikit-learn>=1.3.0numpy>=1.24.0tqdm>=4.66.0accelerate>=0.27.0 ``` Looks harmless. It isn’t. This is the part I want to actually call out, because it’s the difference between a 10-minute job and a 45-minute job spent staring at a CUDA error. **Pin your dependencies to your GPU, not just your framework.** torch>=2.0.0 will happily resolve to a wheel that either doesn't have a CUDA build matching your driver, or pulls a CUDA version your GPU architecture doesn't support well. An RTX Pro 6000 (Blackwell-class) wants a recent CUDA build behind PyTorch, something that actually ships kernels for its compute capability. Throw that same loose torch>=2.0.0 at an older L40S or a V100 and you can end up with a mismatched build, silent CPU fallback, or a job that OOMs in a way that has nothing to do with your batch size. So before you fire off a job: This is exactly the kind of thing that eats people’s time when they’re moving fast, and it’s the reason I built the container/environment layer in Crunr the way I did: controlled, versioned images instead of “pip install and pray” on a fresh box every time. Here’s the full script, train.py. Nothing fancy: standard HuggingFace training loop, tokenization, a base-model baseline pass, per-epoch evaluation, best-checkpoint saving, and a formatted comparison report at the end. ``` """Fine-tune DistilBERT on AG News dataset for topic classification. php AG News label mapping: 0 -> World 1 -> Sports 2 -> Business 3 -> Sci/Tech Usage: python train.py # full fine-tune + eval python train.py --eval-only # only run evaluation (needs saved model) python train.py --epochs 5 # custom epochs""" python import osimport timeimport argparseimport warningswarnings.filterwarnings("ignore") python import numpy as npimport torchfrom torch.utils.data import DataLoaderfrom torch.optim import AdamW python from transformers import ( DistilBertTokenizerFast, DistilBertForSequenceClassification, get_linear_schedule_with_warmup,)from datasets import load_datasetfrom sklearn.metrics import ( accuracy_score, f1_score, precision_score, recall_score, classification_report,)from tqdm import tqdm # ConfigBASE_MODEL = "distilbert-base-uncased"OUTPUT_DIR = "./outputs/finetuned_model"RESULTS_FILE = "./outputs/evaluation_results.txt"MAX_LEN = 128BATCH_SIZE = 32EPOCHS = 3LR = 2e-5WARMUP_RATIO = 0.1SEED = 42NUM_LABELS = 4 LABEL_NAMES = ["World", "Sports", "Business", "Sci/Tech"] python def set_seed(seed: int): torch.manual_seed(seed) np.random.seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) python def load_ag_news(tokenizer, max_len: int, batch_size: int): """Download AG News and return (train_loader, test_loader).""" print("\n[1/5] Loading AG News dataset ...") dataset = load_dataset("fancyzhx/ag_news") python def tokenize(batch): return tokenizer( batch["text"], padding="max_length", truncation=True, max_length=max_len, ) tokenized = dataset.map(tokenize, batched=True, batch_size=1000) tokenized.set_format( type="torch", columns=["input_ids", "attention_mask", "label"], ) train_loader = DataLoader( tokenized["train"], batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True ) test_loader = DataLoader( tokenized["test"], batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True ) print(f" Train samples : {len(dataset['train']):,}") print(f" Test samples : {len(dataset['test']):,}") return train_loader, test_loader php def evaluate(model, loader, device) -> dict: """Run inference and return metric dict.""" model.eval() all_preds, all_labels = [], [] with torch.no_grad(): for batch in tqdm(loader, desc=" Evaluating", leave=False): input_ids = batch["input_ids"].to(device) attention_mask = batch["attention_mask"].to(device) labels = batch["label"].to(device) outputs = model(input_ids=input_ids, attention_mask=attention_mask) preds = torch.argmax(outputs.logits, dim=-1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) metrics = { "accuracy" : accuracy_score(all_labels, all_preds), "f1_macro" : f1_score(all_labels, all_preds, average="macro"), "f1_weighted": f1_score(all_labels, all_preds, average="weighted"), "precision_macro": precision_score(all_labels, all_preds, average="macro"), "recall_macro" : recall_score(all_labels, all_preds, average="macro"), "report" : classification_report( all_labels, all_preds, target_names=LABEL_NAMES, digits=4 ), } return metrics php def evaluate_base_model(tokenizer, test_loader, device) -> dict: """ Evaluate the raw pretrained DistilBERT (randomly initialized classification head) to establish a baseline. """ print("\n[2/5] Evaluating BASE model (no fine-tuning) ...") base_model = DistilBertForSequenceClassification.from_pretrained( BASE_MODEL, num_labels=NUM_LABELS ).to(device) metrics = evaluate(base_model, test_loader, device) del base_model if torch.cuda.is_available(): torch.cuda.empty_cache() return metrics php def train(model, train_loader, test_loader, device, epochs: int) -> dict: """Fine-tune and return best test metrics.""" print(f"\n[3/5] Fine-tuning for {epochs} epoch(s) ...") total_steps = len(train_loader) * epochs warmup_steps = int(total_steps * WARMUP_RATIO) optimizer = AdamW(model.parameters(), lr=LR, weight_decay=0.01) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps, ) best_f1 = 0.0 best_metrics = {} train_start = time.time() for epoch in range(1, epochs + 1): model.train() epoch_loss = 0.0 epoch_start = time.time() pbar = tqdm(train_loader, desc=f" Epoch {epoch}/{epochs}", leave=True) for step, batch in enumerate(pbar, 1): input_ids = batch["input_ids"].to(device) attention_mask = batch["attention_mask"].to(device) labels = batch["label"].to(device) optimizer.zero_grad() outputs = model( input_ids=input_ids, attention_mask=attention_mask, labels=labels, ) loss = outputs.loss loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() epoch_loss += loss.item() pbar.set_postfix(loss=f"{epoch_loss / step:.4f}") epoch_time = time.time() - epoch_start avg_loss = epoch_loss / len(train_loader) print(f" Epoch {epoch}, avg loss: {avg_loss:.4f} ({epoch_time:.1f}s)") # per-epoch evaluation metrics = evaluate(model, test_loader, device) print( f" Val acc={metrics['accuracy']:.4f} " f"f1_macro={metrics['f1_macro']:.4f} " f"precision={metrics['precision_macro']:.4f} " f"recall={metrics['recall_macro']:.4f}" ) if metrics["f1_macro"] > best_f1: best_f1 = metrics["f1_macro"] best_metrics = metrics model.save_pretrained(OUTPUT_DIR) print(f" New best model saved (f1={best_f1:.4f})") total_time = time.time() - train_start print(f"\n Training complete in {total_time / 60:.1f} min") return best_metrics python def write_results(base_metrics: dict, ft_metrics: dict, path: str): """Write side-by-side comparison to a .txt file.""" sep = "=" * 70 sep2 = "-" * 70 python def delta(ft_val, base_val): d = ft_val - base_val sign = "+" if d >= 0 else "" return f"{sign}{d:.4f}" lines = [ sep, " DistilBERT - AG News Topic Classification - Evaluation Report", sep, "", f" Base model : {BASE_MODEL}", f" Fine-tuned epochs: {EPOCHS}", f" Max sequence len : {MAX_LEN}", f" Batch size : {BATCH_SIZE}", f" Learning rate : {LR}", "", sep2, " SUMMARY - Test-set metrics", sep2, f" {'Metric':<25} {'Base':>10} {'Fine-tuned':>12} {'Delta':>10}", sep2, f" {'Accuracy':<25} {base_metrics['accuracy']:>10.4f} {ft_metrics['accuracy']:>12.4f} {delta(ft_metrics['accuracy'], base_metrics['accuracy']):>10}", f" {'F1 (macro)':<25} {base_metrics['f1_macro']:>10.4f} {ft_metrics['f1_macro']:>12.4f} {delta(ft_metrics['f1_macro'], base_metrics['f1_macro']):>10}", f" {'F1 (weighted)':<25} {base_metrics['f1_weighted']:>10.4f} {ft_metrics['f1_weighted']:>12.4f} {delta(ft_metrics['f1_weighted'], base_metrics['f1_weighted']):>10}", f" {'Precision (macro)':<25} {base_metrics['precision_macro']:>10.4f} {ft_metrics['precision_macro']:>12.4f} {delta(ft_metrics['precision_macro'], base_metrics['precision_macro']):>10}", f" {'Recall (macro)':<25} {base_metrics['recall_macro']:>10.4f} {ft_metrics['recall_macro']:>12.4f} {delta(ft_metrics['recall_macro'], base_metrics['recall_macro']):>10}", sep2, "", sep2, " BASE MODEL - Per-class Classification Report", sep2, base_metrics["report"], "", sep2, " FINE-TUNED MODEL - Per-class Classification Report", sep2, ft_metrics["report"], "", sep, " End of report", sep, ] os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True) with open(path, "w") as f: f.write("\n".join(lines)) php print(f"\n[5/5] Results saved -> {path}") python def parse_args(): p = argparse.ArgumentParser(description="Fine-tune DistilBERT on AG News") p.add_argument("--epochs", type=int, default=EPOCHS, help="Training epochs") p.add_argument("--batch-size",type=int, default=BATCH_SIZE, help="Batch size") p.add_argument("--lr", type=float, default=LR, help="Learning rate") p.add_argument("--max-len", type=int, default=MAX_LEN, help="Max token length") p.add_argument("--output-dir",type=str, default=OUTPUT_DIR, help="Where to save model") p.add_argument("--results", type=str, default=RESULTS_FILE,help="Results .txt path") p.add_argument("--eval-only", action="store_true", help="Skip training; load saved model and re-evaluate") return p.parse_args() python def main(): args = parse_args() global EPOCHS, BATCH_SIZE, LR, MAX_LEN, OUTPUT_DIR, RESULTS_FILE EPOCHS = args.epochs BATCH_SIZE = args.batch_size LR = args.lr MAX_LEN = args.max_len OUTPUT_DIR = args.output_dir RESULTS_FILE= args.results set_seed(SEED) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"\n Device : {device}") if torch.cuda.is_available(): print(f" GPU : {torch.cuda.get_device_name(0)}") print(f"\n[0/5] Loading tokenizer ({BASE_MODEL}) ...") tokenizer = DistilBertTokenizerFast.from_pretrained(BASE_MODEL) train_loader, test_loader = load_ag_news(tokenizer, MAX_LEN, BATCH_SIZE) base_metrics = evaluate_base_model(tokenizer, test_loader, device) print( f" Base acc={base_metrics['accuracy']:.4f} " f"f1_macro={base_metrics['f1_macro']:.4f} " f"precision={base_metrics['precision_macro']:.4f} " f"recall={base_metrics['recall_macro']:.4f}" ) if args.eval_only: print(f"\n[3/5] --eval-only: loading fine-tuned model from {OUTPUT_DIR} ...") if not os.path.isdir(OUTPUT_DIR): raise FileNotFoundError( f"No saved model found at '{OUTPUT_DIR}'. Run without --eval-only first." ) ft_model = DistilBertForSequenceClassification.from_pretrained(OUTPUT_DIR).to(device) print("\n[4/5] Evaluating fine-tuned model ...") ft_metrics = evaluate(ft_model, test_loader, device) else: ft_model = DistilBertForSequenceClassification.from_pretrained( BASE_MODEL, num_labels=NUM_LABELS ).to(device) ft_metrics = train(ft_model, train_loader, test_loader, device, EPOCHS) print("\n[4/5] Final evaluation with best checkpoint ...") ft_model = DistilBertForSequenceClassification.from_pretrained(OUTPUT_DIR).to(device) ft_metrics = evaluate(ft_model, test_loader, device) print( f" Fine-tuned acc={ft_metrics['accuracy']:.4f} " f"f1_macro={ft_metrics['f1_macro']:.4f} " f"precision={ft_metrics['precision_macro']:.4f} " f"recall={ft_metrics['recall_macro']:.4f}" ) write_results(base_metrics, ft_metrics, RESULTS_FILE) print("\n" + "=" * 70) print(" FINAL COMPARISON") print("=" * 70) header = f" {'Metric':<25} {'Base':>10} {'Fine-tuned':>12} {'Delta':>10}" print(header) print("-" * 70) for name, key in [ ("Accuracy", "accuracy"), ("F1 (macro)", "f1_macro"), ("F1 (weighted)", "f1_weighted"), ("Precision (macro)", "precision_macro"), ("Recall (macro)", "recall_macro"), ]: b = base_metrics[key] f = ft_metrics[key] d = f - b sign = "+" if d >= 0 else "" print(f" {name:<25} {b:>10.4f} {f:>12.4f} {sign}{d:>9.4f}") print("=" * 70) if __name__ == "__main__": main() ``` Once the environment’s sorted, this is the entire “infra” step: ``` pip install crunr ``` Grab an API key from [cloud.crunr.com](https://cloud.crunr.com/), then: ``` crunr run train.py --instance rtxpro6000 ``` That’s it. One line. Crunr provisions the instance, ships the code, runs it, streams the logs back, and tears the instance down when it’s done, so you’re not paying for idle GPU time while you’re reading the eval report. The RTX Pro 6000 is genuinely overkill for a DistilBERT job this size, you could run this on something far cheaper, but the point of the exercise was showing the ceiling, not the floor. Pick your instance size to match your model and dataset, not your ego. Base model (no fine-tuning, just DistilBERT with a randomly initialized classification head sitting on top): Metric Base Fine-tuned Delta Accuracy 0.2350 0.9455 +0.7105 F1 (macro) 0.1303 0.9456 +0.8153 F1 (weighted) 0.1303 0.9456 +0.8153 Precision (macro) 0.3266 0.9456 +0.6190 Recall (macro) 0.2350 0.9455 +0.7105 The base model’s behavior is a good little lesson on its own. It wasn’t randomly guessing across classes, it was collapsing almost everything into “Sci/Tech” (85.8% recall on that class alone, near-zero on the others). That’s what an untrained head on top of a good encoder looks like: the representations are fine, the classifier on top is noise, and it finds whatever local minimum requires the least commitment. After 3 epochs: Class Precision Recall F1 World 0.9602 0.9532 0.9567 Sports 0.9873 0.9853 0.9863 Business 0.9176 0.9142 0.9159 Sci/Tech 0.9174 0.9295 0.9234 Sports is the easy class here, it has the most distinctive vocabulary, and Business/Sci-Tech are the ones that bleed into each other a little, which tracks with how those two categories actually overlap in real news writing (a chip earnings report is both). Not to flex a 94.5% F1 on a benchmark that’s been solved a hundred times over. The point was the loop: cold GPU, environment that just works, training, honest evaluation, report, all in single-digit minutes, without babysitting a provisioning dashboard or debugging a CUDA mismatch at minute 25 of your 10. That loop is the actual product. The model is just proof it ran. If you want to poke at the script yourself, or try it against a different GPU tier, the whole thing’s in the video, and Crunr’s free to try at [cloud.crunr.com](https://cloud.crunr.com/). I’d genuinely like to hear where it breaks for you, that’s usually more useful to me than someone telling me it worked. [I Have 10 Minutes to Train an AI Model. Here’s Exactly What Happened.](https://blog.stackademic.com/i-have-10-minutes-to-train-an-ai-model-heres-exactly-what-happened-9c1fa24ac78b) was originally published in [Stackademic](https://blog.stackademic.com) on Medium, where people are continuing the conversation by highlighting and responding to this story.