Fix SpecAugment torch.compile graph breaks, deduplicate helpers, adaptive num_workers

[Vaibhavdixit02]

Summary

• SpecAugment rewrite: Replace .item() calls with masked_fill_ and on-device randint so torch.compile can trace the full forward pass without graph breaks (benefits all GPUs, especially A100/H100)
• Deduplicate load_model() and get_device(): Shared helpers in model.py, removing duplicate implementations from transcribe.py and live.py. Consistent cuda > mps > cpu device selection everywhere (previously train.py/eval.py skipped MPS)
• Adaptive num_workers: Default to min(4, os.cpu_count()) instead of hardcoded 4, avoiding the "excessive worker creation" warning on 2-CPU Colab free tier
• Add nanoasr/__init__.py: Public API exports (Conformer, get_config, load_model, evaluate_checkpoint, etc.)
• Fix notebook metadata: Cell 11 (markdown) had invalid execution_count/outputs fields; cell 12 (code) was missing them

Test plan

• All 36 tests pass (pytest tests/ -v)
• from nanoasr import Conformer, load_model, get_device works
• Verify no torch.compile graph break warnings on Colab GPU runtime
• Verify no num_workers warning on 2-CPU Colab free tier

Made with Cursor

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB