Set correct device for all conds

2024-04-18 03:14:28 +05:00
parent 18fec06edb
commit 816cf32335
1 changed files with 1 additions and 1 deletions
@@ -17,7 +17,7 @@ def copy_cond(cond: list[tuple[torch.Tensor, dict]], device="cpu") -> list[list[

 def cat_cond(clip, *conds: list[list[torch.Tensor | dict]], device="cpu") -> list[list[torch.Tensor | dict]]:
    output = copy_cond(conds[0], device=device)
-    cond_tensors = list(map(lambda c: c[0][0], conds))
+    cond_tensors = list(map(lambda c: c[0][0].to(device=device), conds))

    chunk_size = 77
    chunks = list(map(lambda c: c.shape[1] // chunk_size, cond_tensors))