#! /bin/bash
# Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved.

NV_DOCKER=${NV_DOCKER:-"docker"}

DOCKER_ARGS=""
NV_DOCKER_ARGS=""
while [ $# -gt 0 ]; do
    arg=$1
    shift
    DOCKER_ARGS="$DOCKER_ARGS $arg"
    case $arg in
        run|create)
            NV_DOCKER_ARGS="--runtime=nvidia"
            if [ ! -z $NV_GPU ]; then
                NV_DOCKER_ARGS="$NV_DOCKER_ARGS -e NVIDIA_VISIBLE_DEVICES=${NV_GPU// /,}"
            fi
            break
        ;;
        version)
            printf "NVIDIA Docker: 2.0.2\n"
            break
        ;;
        --)
            break
        ;;
    esac
done

if [ ! -z $NV_DEBUG ]; then
    set -x
fi

$NV_DOCKER $DOCKER_ARGS $NV_DOCKER_ARGS "$@"
