Problem when training with 4090#49

Hi, thanks for the great work. I'm trying to run the training code, but when I do pre_train with multiple 4090 gpus, it always gets stuck and no report is shown. But when training with multiple 3090 and single 4090 everything is fine. I strongly suspect that a deadlock occurs in 4090. I narrow the problem down to `deepspeed.initialize` in agent.py. But I don't know how to solve this problem. Any response will be greatly appreciated.