From ec9dedb5f04c35112e49a1f4649e84cca46c22d8 Mon Sep 17 00:00:00 2001 From: Andrew Kvochko Date: Sun, 24 Nov 2024 11:53:12 +0000 Subject: [PATCH] Revert an unnecessary change. --- t5_encoder.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/t5_encoder.py b/t5_encoder.py index 1c62c89..9ba6542 100644 --- a/t5_encoder.py +++ b/t5_encoder.py @@ -78,13 +78,13 @@ class LTXVTextEncoderModel(torch.nn.Module): token_weight_pairs_t5 = token_weight_pairs["t5xxl"] text_input_ids = torch.tensor( [[t[0] for t in token_weight_pairs_t5]], - device=self.device, + device=self.t5xxl.device, ) prompt_attention_mask = torch.tensor( [[w[1] for w in token_weight_pairs_t5]], - device=self.device, + device=self.t5xxl.device, ) - self.to(self.device) # comfyui skips loading some weights to gpu + self.to(self.t5xxl.device) # comfyui skips loading some weights to gpu out = self.t5xxl(text_input_ids, attention_mask=prompt_attention_mask)[0] out = out * prompt_attention_mask.unsqueeze(2) return out, None, {"attention_mask": prompt_attention_mask}