Skip to content

GitLab

Commit 909d4f1a authored by thomwolf's avatar thomwolf
Browse files

cuda again

parent 14f0e8e5
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
examples/bertology.py
View file @ 909d4f1a
...@@ -218,7 +218,7 @@ def run_model(): ...@@ -218,7 +218,7 @@ def run_model():
print_2d_tensor(head_importance) print_2d_tensor(head_importance)
logger.info("Head ranked by importance scores") logger.info("Head ranked by importance scores")
head_ranks = torch.zeros(head_importance.numel(), dtype=torch.long, device=args.device) head_ranks = torch.zeros(head_importance.numel(), dtype=torch.long, device=args.device)
head_ranks[head_importance.view(-1).sort(descending=True)[1]] = torch.arange(head_importance.numel()) head_ranks[head_importance.view(-1).sort(descending=True)[1]] = torch.arange(head_importance.numel(), device=args.device)
head_ranks = head_ranks.view_as(head_importance) head_ranks = head_ranks.view_as(head_importance)
print_2d_tensor(head_ranks) print_2d_tensor(head_ranks)
......
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