daking vchiley commited on
Commit
e837ad7
1 Parent(s): 305ddb5

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

Browse files

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


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

Files changed (1) hide show
  1. attention.py +1 -1
attention.py CHANGED
@@ -55,7 +55,7 @@ def scaled_multihead_dot_product_attention(query, key, value, n_heads, past_key_
55
  attn_weight = torch.softmax(attn_weight, dim=-1)
56
  if dropout_p:
57
  attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p, training=training, inplace=True)
58
- out = attn_weight.matmul(v)
59
  out = rearrange(out, 'b h s d -> b s (h d)')
60
  if needs_weights:
61
  return (out, attn_weight, past_key_value)
 
55
  attn_weight = torch.softmax(attn_weight, dim=-1)
56
  if dropout_p:
57
  attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p, training=training, inplace=True)
58
+ out = attn_weight.to(v.dtype).matmul(v)
59
  out = rearrange(out, 'b h s d -> b s (h d)')
60
  if needs_weights:
61
  return (out, attn_weight, past_key_value)