{

py=$1

device=$2

tensorflow_pkg='tensorflow'

if [[ ${device} == 'gpu' ]]; then

tensorflow_pkg='tensorflow-gpu'

fi

tag=horovod-build-py${py}-${device}:$(date +%Y%m%d-%H%M%S)

docker build -f Dockerfile.${device} -t ${tag} --build-arg python=${py} --no-cache .

horovod_version=$(docker run --rm ${tag} pip freeze | grep ^horovod= | awk -F== '{print $2}')

tensorflow_version=$(docker run --rm ${tag} pip freeze | grep ^${tensorflow_pkg}= | awk -F== '{print $2}')

pytorch_version=$(docker run --rm ${tag} pip freeze | grep ^torch= | awk -F== '{print $2}' | awk -F+ '{print $1}')

mxnet_version=$(docker run --rm ${tag} pip freeze | grep ^mxnet | awk -F== '{print $2}')

final_tag=horovod/horovod:${horovod_version}-tf${tensorflow_version}-torch${pytorch_version}-mxnet${mxnet_version}-py${py}-${device}

docker tag ${tag} ${final_tag}

docker rmi ${tag}