Thanks for the reply!  I use a colab session.
restarted the sesion between faillures.
I have in the code the following line:
!export CUDA_LAUNCH_BLOCKING=1
and I have two errors depending if I have a print or not. The problem I think is with the logit tensor.
The error is:
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-43-9064aecd8313> in <module>()
     18         #print("[evaluation] tmp_eval_loss: ", tmp_eval_loss.shape)
     19         logits = model(b_input_ids, token_type_ids=None,
---> 20                        attention_mask=b_input_mask)[0]
     21         #print("logits: ", logits.shape)
     22         #print("logits: ", logits.detach().cpu().numpy())

5 frames
/usr/local/lib/python3.6/dist-packages/torch/_tensor_str.py in __init__(self, tensor)
     85 
     86         else:
---> 87             nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
     88 
     89             if nonzero_finite_vals.numel() == 0:

RuntimeError: cuda runtime error (710) : device-side assert triggered at /pytorch/aten/src/THC/THCReduceAll.cuh:327

and the stack
RuntimeError                              Traceback (most recent call last)
<ipython-input-43-9064aecd8313> in <module>()
     18         #print("[evaluation] tmp_eval_loss: ", tmp_eval_loss.shape)
     19         logits = model(b_input_ids, token_type_ids=None,
---> 20                        attention_mask=b_input_mask)[0]
     21         #print("logits: ", logits.shape)
     22         #print("logits: ", logits.detach().cpu().numpy())

5 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    539             result = self._slow_forward(*input, **kwargs)
    540         else:
--> 541             result = self.forward(*input, **kwargs)
    542         for hook in self._forward_hooks.values():
    543             hook_result = hook(self, input, result)

<ipython-input-28-7837bc2054a4> in forward(self, input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, labels)
    196             outputs = (loss,) + outputs
    197         # <------- end of loss 3 ------->
--> 198         print("outputs: ", outputs)
    199         return outputs  # (loss), scores, (hidden_states), (attentions)

/usr/local/lib/python3.6/dist-packages/torch/tensor.py in __repr__(self)
    128         # characters to replace unicode characters with.
    129         if sys.version_info > (3,):
--> 130             return torch._tensor_str._str(self)
    131         else:
    132             if hasattr(sys.stdout, 'encoding'):

/usr/local/lib/python3.6/dist-packages/torch/_tensor_str.py in _str(self)
    309                 tensor_str = _tensor_str(self.to_dense(), indent)
    310             else:
--> 311                 tensor_str = _tensor_str(self, indent)
    312 
    313     if self.layout != torch.strided:

/usr/local/lib/python3.6/dist-packages/torch/_tensor_str.py in _tensor_str(self, indent)
    207     if self.dtype is torch.float16 or self.dtype is torch.bfloat16:
    208         self = self.float()
--> 209     formatter = _Formatter(get_summarized_data(self) if summarize else self)
    210     return _tensor_str_with_formatter(self, indent, formatter, summarize)
    211 

/usr/local/lib/python3.6/dist-packages/torch/_tensor_str.py in __init__(self, tensor)
     85 
     86         else:
---> 87             nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) & tensor_view.ne(0))
     88 
     89             if nonzero_finite_vals.numel() == 0:

RuntimeError: cuda runtime error (710) : device-side assert triggered at /pytorch/aten/src/THC/THCReduceAll.cuh:327

And if I remove the print then I have the following error
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-43-9064aecd8313> in <module>()
     26     #print("it:", it)
     27     #print("[evaluation] logits: ", logits)
---> 28     logits = logits.cpu().data.numpy()
     29     #print("logits: ", logits)
     30 

RuntimeError: CUDA error: device-side assert triggered