...
The last function call sets the GPU Tensorflow will use for each process.
Try using as a template, To test what we have written so far, use the batch job script runTensorflow.sh provided in the previous page as a template for submitting the job. You will need to adapt it for the batch job script and change the number of GPUs per node together with the srun command, and to use of the python script porposed here: (01_horovod_mnist.py) containing the two parts described above. The adapted lines of the batch job script should look like:
#SBATCH --nodes=2 #2 nodes in this example #SBATCH --gres=gpu:2 #2 GPUS per node..PYTHON_SCRIPT=$PYTHON_SCRIPT_DIR/01_horovod_mnist.py..srun -N 2 -n 4 -c 8 --gres=gpu:2 python3 $PYTHON_SCRIPT
...