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

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


Co-authored-by: Vitaliy Chiley <[email protected]>

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)