arcface_torch training subprocess error when config.dali = True

Traceback (most recent call last):
File "train.py", line 151, in
main(parser.parse_args())
File "train.py", line 108, in main
loss: torch.Tensor = module_partial_fc(local_embeddings, local_labels, opt)
File "/train_tmp/miniconda3/envs/dali/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "/data/AI-GPU/chenguote/face_recognition/insightface/recognition/arcface_torch/partial_fc.py", line 243, in forward
distributed.all_gather(_gather_labels, local_labels)
File "/train_tmp/miniconda3/envs/dali/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 1863, in all_gather
work = default_pg.allgather([tensor_list], [tensor])
RuntimeError: All tensor operands to scatter/gather must have the same number of elements
Killing subprocess 2701979
Killing subprocess 2701980
Killing subprocess 2701982
Killing subprocess 2701983
Killing subprocess 2701984
Killing subprocess 2701987
Killing subprocess 2701990
Killing subprocess 2701993
Traceback (most recent call last):
File "/train_tmp/miniconda3/envs/dali/lib/python3.7/runpy.py", line 193, in _run_module_as_main
"__main__", mod_spec)
File "/train_tmp/miniconda3/envs/dali/lib/python3.7/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/train_tmp/miniconda3/envs/dali/lib/python3.7/site-packages/torch/distributed/launch.py", line 340, in
main()
File "/train_tmp/miniconda3/envs/dali/lib/python3.7/site-packages/torch/distributed/launch.py", line 326, in main
sigkill_handler(signal.SIGTERM, None) # not coming back
File "/train_tmp/miniconda3/envs/dali/lib/python3.7/site-packages/torch/distributed/launch.py", line 301, in sigkill_handler
raise subprocess.CalledProcessError(returncode=last_return_code, cmd=cmd)
subprocess.CalledProcessError: Command '['/train_tmp/miniconda3/envs/dali/bin/python', '-u', 'train.py', '--local_rank=7', 'configs/config.py']' returned non-zero exit status 1.

config.dali set False is ok and got 4400 samples/sec