Hi! I am pretty new to Hugging Face and I am struggling with next sentence prediction model. I would like it to use a GPU device inside a Colab Notebook but I am not able to do it. This is my proposal:
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForNextSentencePrediction.from_pretrained('bert-base-uncased', return_dict=True)
model.to("cuda:0")
prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
next_sentence = "The sky is blue due to the shorter wavelength of blue light."
tokenizer_output = tokenizer(prompt, next_sentence, return_tensors='pt')
tokens_tensor = tokenizer_output['input_ids'].to('cuda:0')
token_type_ids = tokenizer_output['token_type_ids'].to('cuda:0')
attention_mask = tokenizer_output['attention_mask'].to('cuda:0')
encoding = {'input_ids' : tokens_tensor, 
        'token_type_ids' : token_type_ids, 
        'attention_mask' : attention_mask}
outputs = model(**encoding)
logits = outputs.logits
print(logits) # next sentence was random

However, it does not work. The GPU device is properly set, since I used the following command to check it:
import torch; print(torch.cuda.get_device_name(0))

I would appreciate any help