diff --git a/.github/workflows/self-push.yml b/.github/workflows/self-push.yml
index fceb9400d1..7a11c401b9 100644
--- a/.github/workflows/self-push.yml
+++ b/.github/workflows/self-push.yml
@@ -36,7 +36,7 @@ jobs:
       run: |
         source .env/bin/activate
         pip install --upgrade pip
-        pip install torch --no-cache-dir
+        pip install torch!=1.6.0 --no-cache-dir
         pip install .[sklearn,testing]
 
     - name: Are GPUs recognized by our DL frameworks
diff --git a/src/transformers/modeling_xlm.py b/src/transformers/modeling_xlm.py
index 27a8ed21aa..26a514b716 100644
--- a/src/transformers/modeling_xlm.py
+++ b/src/transformers/modeling_xlm.py
@@ -496,11 +496,13 @@ class XLMModel(XLMPreTrainedModel):
         else:
             bs, slen = inputs_embeds.size()[:-1]
 
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
         if lengths is None:
             if input_ids is not None:
                 lengths = (input_ids != self.pad_index).sum(dim=1).long()
             else:
-                lengths = torch.LongTensor([slen] * bs)
+                lengths = torch.tensor([slen] * bs, device=device)
         # mask = input_ids != self.pad_index
 
         # check inputs