fix: t5 embedder device

archinetai · Oct 13, 2022 · 7fabeb3 · 7fabeb3
1 parent d3f27f8
commit 7fabeb3
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/audio_diffusion_pytorch/modules.py b/audio_diffusion_pytorch/modules.py
@@ -1153,10 +1153,14 @@ def forward(self, texts: List[str]) -> Tensor:
             return_tensors="pt",
         )
 
+        device = next(self.transformer.parameters()).device
+        input_ids = encoded["input_ids"].to(device)
+        attention_mask = encoded["attention_mask"].to(device)
+
         self.transformer.eval()
 
         embedding = self.transformer(
-            input_ids=encoded["input_ids"], attention_mask=encoded["attention_mask"]
+            input_ids=input_ids, attention_mask=attention_mask
         )["last_hidden_state"]
 
         return embedding

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
     name="audio-diffusion-pytorch",
     packages=find_packages(exclude=[]),
-    version="0.0.63",
+    version="0.0.64",
     license="MIT",
     description="Audio Diffusion - PyTorch",
     long_description_content_type="text/markdown",