A further question on this:
I want to use multiple gpu, so I do:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def get_bert():
    from transformers import BertTokenizer, BertModel
    tokenizer = BertTokenizer.from_pretrained("hfl/chinese-roberta-wwm-ext", model_max_length=128)
    model = BertModel.from_pretrained("hfl/chinese-roberta-wwm-ext")
    if torch.cuda.device_count() > 1:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
        model = torch.nn.DataParallel(model)
        #model = torch.nn.DataParallel(model, device_ids=[0, 1])
    model = model.to(device)

    return tokenizer, model

From nvidia-smi screenshot, you can see that there are 4 GPUs, however, the torch.cuda.device_count()=3, and my program only uses 3 GPUs(0,1,2), while GPU 3 is idle. If I specify:
device_ids=[0,1,2,3]
It says ‘AssertionError: Invalid device id’.
Another question is when using multl-GPUs, is this specification of device right?
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
This seems to only get ‘cuda:0’, but the code above does use 3 GPUs, which I don’t quite understand.