[CheckpointWrapper] only checkpoint the transform method, automatically add .pkl extension to model names