Commit 25d5ca48 authored by VictorSanh's avatar VictorSanh
Browse files

Fix scatter LopngTensor

parent e6a710f6
...@@ -493,7 +493,7 @@ class BertForQuestionAnswering(nn.Module): ...@@ -493,7 +493,7 @@ class BertForQuestionAnswering(nn.Module):
def compute_loss(logits, positions): def compute_loss(logits, positions):
max_position = positions.max().item() 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, device=input_ids.device).zero_()
one_hot = one_hot.scatter(1, positions, 1) 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]
log_probs = nn.functional.log_softmax(logits, dim = -1).view(batch_size, seq_length) 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) loss = -torch.mean(torch.sum(one_hot*log_probs), dim = -1)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment