def train_model(state, train_loader, num_epochs=10):
    for epoch in range(num_epochs):
        for batch in train_loader:
            state = train_step(state, batch)
        print(f"Epoch {epoch + 1} completed")
    return state

        __
        __