Like what we did for training acceleration: https://github.com/Project-MONAI/tutorials/blob/main/acceleration/fast_training_tutorial.ipynb