def METHOD_NAME(state, train_ds, test_ds):
  for epoch in range(1, _NUM_EPOCHS.value+1):
    start_time = time.time()
    state, train_loss, train_accuracy = run_epoch(state, train_ds)
    _, test_loss, test_accuracy = run_epoch(state, test_ds, METHOD_NAME=False)
    print(f"Training set accuracy {train_accuracy}")
    print(f"Training set loss {train_loss}")
    print(f"Test set accuracy {test_accuracy}")
    print(f"Test set loss {test_loss}")
    epoch_time = time.time() - start_time
    print(f"Epoch {epoch} in {epoch_time:0.2f} sec")
  return state