From e6a710f68473da14cf9ec50f9e748cfa01a927e5 Mon Sep 17 00:00:00 2001 From: VictorSanh Date: Fri, 2 Nov 2018 17:54:22 -0400 Subject: [PATCH] device --- modeling_pytorch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modeling_pytorch.py b/modeling_pytorch.py index b227dfeb91..07fb256104 100644 --- a/modeling_pytorch.py +++ b/modeling_pytorch.py @@ -492,7 +492,7 @@ class BertForQuestionAnswering(nn.Module): def compute_loss(logits, positions): max_position = positions.max().item() - one_hot = torch.FloatTensor(batch_size, max(max_position, seq_length) +1).zero_() + one_hot = torch.FloatTensor(batch_size, max(max_position, seq_length) +1, device=input_ids.device).zero_() one_hot = one_hot.scatter(1, positions, 1) one_hot = one_hot[:, :seq_length] log_probs = nn.functional.log_softmax(logits, dim = -1).view(batch_size, seq_length)