To enable models built using PyTorch to use the pipeline for end to end runs. Relates to icenet-ai/icenet-roadmap#54