From cb7811706e6a29bad4af857d060f40bc705d1d46 Mon Sep 17 00:00:00 2001 From: ssit Date: Thu, 7 Sep 2023 11:39:54 -0400 Subject: [PATCH] Ensure attn bias is the same dtype as query --- fabric/weighted_attn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fabric/weighted_attn.py b/fabric/weighted_attn.py index 5434205..bd41aea 100644 --- a/fabric/weighted_attn.py +++ b/fabric/weighted_attn.py @@ -110,5 +110,5 @@ def get_attn_bias(weights, shape, dtype): attn_bias = weights.log().clamp(min=min_val) view_shape = [B, *([1] * (len(shape) - 2)), nk] # print("[Weighted Attn] get_attn_bias:", attn_bias.shape, view_shape, shape) - attn_bias = attn_bias.view(view_shape).expand(shape) + attn_bias = attn_bias.view(view_shape).expand(shape).to(dtype) return attn_bias