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

nlp
1

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_select(b_labels, (preds_mask == 1))
tr_batch_preds = np.argmax(tr_logits[preds_mask.squeeze()], axis=1)
tr_batch_labels = tr_label_ids.to(device).numpy()
tr_preds.extend(tr_batch_preds)
tr_labels.extend(tr_batch_labels)

Error Message:

Using TensorFlow backend.
Saved standardized data to ./data/en/combined/train_combined.txt.
Saved standardized data to ./data/en/combined/dev_combined.txt.
Saved standardized data to ./data/en/combined/test_combined.txt.
Constructed SentenceGetter with 25650 examples.
Constructed SentenceGetter with 8934 examples.
Loaded training and validation data into DataLoaders.
Initialized model and moved it to cuda.
Initialized optimizer and set hyperparameters.
Epoch:   0% 0/5 [00:00<?, ?it/s]Starting training loop.
Epoch:   0% 0/5 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "/content/FYP_Presentation/python/main.py", line 102, in <module>
    valid_dataloader,
  File "/content/FYP_Presentation/python/utils/main_utils.py", line 431, in train_and_save_model
    tr_batch_preds = torch.max(tr_logits[preds_mask.squeeze()], axis=1)
  File "/usr/local/lib/python3.6/dist-packages/torch/tensor.py", line 412, in __array__
    return self.numpy()
TypeError: can't convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
2

CUDA tensors can not be converted to numpy directly. As the error message says please do tensor_name.cpu().numpy() to convert this tensor to numpy