Birchlabs
/

mosaicml-mpt-7b-chat-qlora

Text Generation

text-generation-inference

Model card Files Files and versions Community

Alex Birch commited on May 23, 2023

Commit

ec8ea9d

·

unverified ·

1 Parent(s): 07e555c

prefer NamedTuple

Files changed (1) hide show

attention.py +1 -1

attention.py CHANGED Viewed

@@ -121,7 +121,7 @@ def scaled_multihead_dot_product_attention(
     out = attn_weight.matmul(v)
     out = rearrange(out, 'b h s d -> b s (h d)')
     if needs_weights:
-        return (out, attn_weight)
     return AttnFnOutput(out, None)
 def check_valid_inputs(*tensors, valid_dtypes=[torch.float16, torch.bfloat16]):

     out = attn_weight.matmul(v)
     out = rearrange(out, 'b h s d -> b s (h d)')
     if needs_weights:
+        return AttnFnOutput(out, attn_weight)
     return AttnFnOutput(out, None)
 def check_valid_inputs(*tensors, valid_dtypes=[torch.float16, torch.bfloat16]):