fix for negpip

hako-mikan · Feb 7, 2024 · 59d68e6 · 59d68e6
1 parent b6e987f
commit 59d68e6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/scripts/attention.py b/scripts/attention.py
@@ -20,6 +20,7 @@ def main_forward(module,x,context,mask,divide,isvanilla = False,userpp = False,t
     if negpip:
         conds, contokens = negpip
         context = torch.cat((context,conds),1)
+        context = context.to(x.dtype)
 
     h = module.heads
     if isvanilla: # SBM Ddim / plms have the context split ahead along with x.
@@ -596,4 +597,4 @@ def negpipdealer(i,pn):
         else:
             return None
     else:
-        return None
+        return None