Regions: also check if dtype matches

2024-09-12 10:17:37 +02:00
parent f986f6a442
commit 61fa161c34
1 changed files with 2 additions and 2 deletions
@@ -150,9 +150,9 @@ class AttentionMask:
            assert k.mean() == v.mean(), "k and v must be the same."
            device, dtype = q.device, q.dtype

-            if self.conds[0].device != device:
+            if self.conds[0].device != device or self.conds[0].dtype != dtype:
                self.conds = [cond.to(device, dtype=dtype) for cond in self.conds]
-            if self.mask.device != device:
+            if self.mask.device != device or self.mask.dtype != dtype:
                self.mask = self.mask.to(device, dtype=dtype)

            cond_or_unconds = extra_options["cond_or_uncond"]