Error when training with multiple GPUs in TAO

Hi,
Could you pull the 4.0.1 version of TAO docker and retry?

$ docker pull nvcr.io/nvidia/tao/tao-toolkit:4.0.1-tf1.15.5

Then, login in the docker
$ docker run --runtime=nvidia -it --rm nvcr.io/nvidia/tao/tao-toolkit:4.0.1-tf1.15.5 /bin/bash

Then, run the commands. Note, “tao” is not needed when run inside the docker.

yolo_v4 train -e $SPECS_DIR/yolo_v4_train_resnet18_kitti.txt
-r $USER_EXPERIMENT_DIR/experiment_dir_unpruned
-k $KEY
--gpus 8