address #15

Author: lucidrains

perfusion_pytorch/embedding.py

@@ -113,6 +113,10 @@ def __init__(
     def parameters(self):
         return [self.concepts]
 
+    @property
+    def device(self):
+        return self.concepts.device
+
     @beartype
     def forward(
         self,
@@ -138,6 +142,7 @@ def forward(
             inferred_concept_id = self.concept_embed_ids[0]
 
             x = self.tokenize(x)
+            x = x.to(self.device)
 
             superclass_mask = x == self.superclass_embed_id
             assert superclass_mask.any(dim = -1).all(), 'superclass embed id must be present for all prompts'

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'perfusion-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.20',
+  version = '0.1.21',
   license='MIT',
   description = 'Perfusion - Pytorch',
   author = 'Phil Wang',