BERT NER: can't convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first

Question

I want to train my BERT NER model on colab. But following error occurs

Code:

tr_logits = tr_logits.detach().cpu().numpy()
tr_label_ids = torch.masked_s