PyTorch: Build a Complete Training Loop