learning_rate = 0.0001
n_epochs = 2000

print("\nTraining...")
for epoch in range(n_epochs):
    # This single line trains ALL THREE models simultaneously
    params, losses = parallel_train_step(params, X, Y, learning_rate)
    
    if epoch % 500 == 0:
        print(f"Epoch {epoch:4d} | Losses: {losses}")

print("\nFinal learned parameters:")
print(f"{'City':<6} {'True Slope':<12} {'Learned Slope':<15} {'True Intercept':<15} {'Learned Intercept':<15}")
print("-" * 70)
for i in range(n_cities):
    print(f"{i:<6} {true_slopes[i]:<12.1f} {params[i, 0]:<15.2f} {true_intercepts[i]:<15.1f} {params[i, 1]:<15.2f}")

        __
        __