diff --git a/examples/bertology.py b/examples/bertology.py
index 888d95e5..6997b9e2 100644
--- a/examples/bertology.py
+++ b/examples/bertology.py
@@ -209,8 +209,8 @@ def run_model():
     attn_entropy, head_importance, _, _ = compute_heads_importance(args, model, eval_dataloader)
 
     # Print/save matrices
-    np.save(os.path.join(args.output_dir, 'attn_entropy.npy'), attn_entropy)
-    np.save(os.path.join(args.output_dir, 'head_importance.npy'), head_importance)
+    np.save(os.path.join(args.output_dir, 'attn_entropy.npy'), attn_entropy.detach().cpu().numpy())
+    np.save(os.path.join(args.output_dir, 'head_importance.npy'), head_importance.detach().cpu().numpy())
 
     logger.info("Attention entropies")
     print_2d_tensor(attn_entropy)