GPU-enabled training loop with basic profiling

Overview: what you’re building

You’re going to build a train_gpu.py script that:

• Uses GPU automatically when available.
• Ensures model, inputs, and loss are on the same device.
• Logs tokens per second so you can see training speed.
• Supports gradient accumulation to simulate larger batch sizes.
• Saves checkpoints of model + optimizer so you can resume training.

We’ll assume PyTorch, a text model (like a small Transformer), and that you already have a basic CPU training loop.