Scaling PyTorch Model Training With Minimal Code Changes