From 816cf3233546b01513b0531175bcc6c6a8cbced1 Mon Sep 17 00:00:00 2001 From: Pam Date: Thu, 18 Apr 2024 03:14:28 +0500 Subject: [PATCH] Set correct device for all conds --- utils/cond_functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/cond_functions.py b/utils/cond_functions.py index efcd38c..f3936c6 100644 --- a/utils/cond_functions.py +++ b/utils/cond_functions.py @@ -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))