Besides, when I run torch.cuda.device_count(), I got “4”, so there should be no problem with driver
torch.cuda.device_count()