diff --git a/examples/bertology.py b/examples/bertology.py
index 6997b9e2..69432816 100644
--- a/examples/bertology.py
+++ b/examples/bertology.py
@@ -218,7 +218,7 @@ def run_model():
     print_2d_tensor(head_importance)
     logger.info("Head ranked by importance scores")
     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)
     print_2d_tensor(head_ranks)