Copy attention tensor to CPU for streamlit caching

Streamlit tries to cache it using numpy (I think), and
this does not work because the tensor is on the GPU.
Interestingly it only causes issues for ZymCTRL not for
TAPE-BERT.

hexviz/attention.py

@@ -89,7 +89,8 @@ def get_attention(
     else:
         raise ValueError(f"Model {model_type} not supported")
 
-    return attentions
+    # Transfer to CPU to avoid issues with streamlit caching
+    return attentions.cpu()
 
 def unidirectional_avg_filtered(attention, layer, head, threshold):
     num_layers, num_heads, seq_len, _ = attention.shape