Fix: copy dtype and device in quat_action allocation

naver · Nov 20, 2023 · 40a0a9f · 40a0a9f
1 parent ef10b1e
commit 40a0a9f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/roma/utils.py b/roma/utils.py
@@ -256,7 +256,7 @@ def quat_action(q, v, is_normalized=False):
     """
     batch_shape = v.shape[:-1]
     iquat = quat_conjugation(q) if is_normalized else quat_inverse(q)
-    pure = torch.cat((v, torch.zeros(batch_shape + (1,))), dim=-1)
+    pure = torch.cat((v, torch.zeros(batch_shape + (1,), dtype=q.dtype, device=q.device)), dim=-1)
     res = quat_product(q, quat_product(pure, iquat))
     return res[...,:3]