def METHOD_NAME(train_dataset, test_dataset, model, epochs, batch_size, lr,
        lr_decay_factor, lr_decay_step_size, weight_decay, inference,
        profiling, bf16, use_compile):
    if not inference:
        run_train(train_dataset, test_dataset, model, epochs, batch_size,
                  use_compile, lr, lr_decay_factor, lr_decay_step_size,
                  weight_decay)
    else:
        run_inference(test_dataset, model, epochs, batch_size, profiling, bf16,
                      use_compile)