From 72ab10399f6b2d6e338aca4ac0ebb2f16a0cfb2e Mon Sep 17 00:00:00 2001
From: VictorSanh <victorsanh@gmail.com>
Date: Fri, 2 Nov 2018 18:06:21 -0400
Subject: [PATCH] Fix loss

Please review @thomwolf but i think this is equivqlent (and it mimics the loss computation of the original loss)
---
 modeling_pytorch.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/modeling_pytorch.py b/modeling_pytorch.py
index 32ae35b74f..8d8c1aeed0 100644
--- a/modeling_pytorch.py
+++ b/modeling_pytorch.py
@@ -492,9 +492,9 @@ 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, device=input_ids.device).zero_()
+                one_hot = torch.FloatTensor(batch_size, max(max_position, seq_length) +1).zero_()
                 one_hot = one_hot.scatter(1, positions.cpu(), 1) # Second argument need to be LongTensor and not cuda.LongTensor
-                one_hot = one_hot[:, :seq_length]
+                one_hot = one_hot[:, :seq_length].to(input_ids.device)
                 log_probs = nn.functional.log_softmax(logits, dim = -1).view(batch_size, seq_length)
                 loss = -torch.mean(torch.sum(one_hot*log_probs), dim = -1)
                 return loss