* Modelled closely after the existing torch benchmark script but just adapted slightly for tensorflow