auto-move inputs to same device as robotic transformer

lucidrains · Jan 31, 2024 · 0d051d5 · 0d051d5
1 parent 9a0fdef
commit 0d051d5
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/q_transformer/q_robotic_transformer.py b/q_transformer/q_robotic_transformer.py
@@ -1207,6 +1207,15 @@ def forward(
         cond_drop_prob = 0.,
     ):
 
+        # just auto-move inputs to the same device as robotic transformer
+
+        video = video.to(self.device)
+
+        if exists(actions):
+            actions = actions.to(self.device)
+
+        # encoding state
+
         encoded_state = self.encode_state(
             video = video,
             texts = texts,

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'q-transformer',
   packages = find_packages(exclude=[]),
-  version = '0.1.11',
+  version = '0.1.12',
   license='MIT',
   description = 'Q-Transformer',
   author = 'Phil Wang',