mirror of
https://github.com/huggingface/transformers.git
synced 2025-08-02 11:11:05 +06:00
device
This commit is contained in:
parent
3ebf1a13c9
commit
e6a710f684
@ -492,7 +492,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).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, 1)
|
||||||
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)
|
||||||
|
Loading…
Reference in New Issue
Block a user