RuntimeError: CUDA error: device-side assert triggered - When calling a model for the second time

I have the following error when using a PyTorch model :

/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)
   2197         # remove once script supports set_grad_enabled
   2198         _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)
-> 2199     return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
   2200 
   2201 

RuntimeError: CUDA error: device-side assert triggered

The error seems to happen only the second time I call the model My code :

epochs =  500
losses = []
model.to(device)

for e in range(epochs):
  running_loss = 0
  current_batch = 1

  for x1, x2, y in data_loader:    
    print("x1 to device")
    x3 = x1.to(device)
    print("--- Computing embedding1 ---")
    embedding1 = model(x3, pooling_method=pooling_method)
    print(embedding1.size())

    print("x2 to device")
    x4 = x2.to(device)
    print("--- Computing embedding2 ---")
    embedding2 = model(x4, pooling_method=pooling_method)
    print(embedding2.size())

The output :

x1 to device
--- Computing embedding1 ---
torch.Size([64, 768])
x2 to device
--- Computing embedding2 ---
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-29-6b36cff704b2> in <module>
     21     x4 = x2.to(device)
     22     print("--- Computing embedding2 ---")
---> 23     embedding2 = model(x4, pooling_method=pooling_method)
     24     print(embedding2.size())
     25 

8 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)
   2197         # remove once script supports set_grad_enabled
   2198         _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)
-> 2199     return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
   2200 
   2201 

RuntimeError: CUDA error: device-side assert triggered

The inputs have the same shape so the problem is not about the shapes. The error seems to happen when the model computes the output, but only the second time.

The device is :

device(type='cuda', index=0)

And if necessary, the model is :

class BERT(nn.Module):
    """
    Torch model based on CamemBERT, in order to make sentence embeddings
    """
    def __init__(self, tokenizer, model_name=model_name, output_size=100):
        super().__init__()

        self.bert = CamembertModel.from_pretrained(model_name)
        self.bert.resize_token_embeddings(len(tokenizer))

        
    def forward(self, x, pooling_method='cls'):
        hidden_states = self.bert(x).last_hidden_state
        embedding = pooling(hidden_states, pooling_method=pooling_method)

        return embedding

Does anyone know how to resolve this ?