start_time = time.time()
torch.manual_seed(123)

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=5e-5,
    weight_decay=0.1
)

num_epochs = 2

train_loss, val_loss, tokens_seen = train_model_simple(
    model,train_loader, val_loader, optimizer, device,
    num_epochs=num_epochs,eval_freq=5,eval_iter=5,
    start_context=format_input(val_data[0]),
    tokenizer=tokenizer
)

end_time = time.time()
execution_time_minutes = (end_time - start_time) / 60
print(f"Fine-tuning completed in {execution_time_minutes:.2f} minutes.")

        __
        __