add explicit cast where running without autocast causes issues (#41)

- add explicit cast where running without autocast causes issues (d35f0e62c295bacd6c29796e453a6f65f4d58ac8)


Co-authored-by: Vitaliy Chiley <vchiley@users.noreply.huggingface.co>

attention.py

@@ -55,7 +55,7 @@ def scaled_multihead_dot_product_attention(query, key, value, n_heads, past_key_
     attn_weight = torch.softmax(attn_weight, dim=-1)
     if dropout_p:
         attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p, training=training, inplace=True)
-    out = attn_weight.matmul(v)
+    out = attn_weight.to(v.dtype).matmul(v)
     out = rearrange(out, 'b h s d -> b s (h d)')
     if needs_weights:
         return (out, attn_weight, past_key_value)