From 0d9328f2ef01d9ddea58b04fd52a013e3c2fd136 Mon Sep 17 00:00:00 2001 From: Lysandre Debut Date: Fri, 7 Aug 2020 02:58:15 -0400 Subject: [PATCH] Patch GPU failures (#6281) * Pin to 1.5.0 * Patch XLM GPU test --- .github/workflows/self-push.yml | 2 +- src/transformers/modeling_xlm.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/self-push.yml b/.github/workflows/self-push.yml index fceb9400d..7a11c401b 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 27a8ed21a..26a514b71 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