def METHOD_NAME(args):
    data_module = MNISTDataModule(
        args.training_data, args.validation_data, args.batch_size, args.shuffle
    )
    # train
    model = LightningMNISTClassifier()
    checkpoint_callback = pl.callbacks.ModelCheckpoint(
        save_top_k=1,
        verbose=True,
        every_n_train_steps=1000,
        monitor="train_loss",
        mode="min",
        save_last=True,
    )
    callbacks = [UpdateDataStepCallback(), checkpoint_callback]
    trainer = pl.Trainer(
        max_epochs=2,
        default_root_dir="./data/",
        callbacks=callbacks,
    )
    trainer.fit(model, data_module)