Adapt optimus to Distributon

This commit adapts the existing code to use the distributed library via a config option.
To achieve this we switch to using Pytorch's dataloader.