Adapted from https://github.com/facebookresearch/mae/tree/main to run with pytorch lightning;
train / fit: python main.py fit --config config.yaml
predict: python main.py predict --config config.yaml
- predict: pl_model, callback, datamodule
- test: pl_model, callback, datamodule