Initialize V_inv on device

eigenvivek · Dec 5, 2023 · 3962fd9 · 3962fd9
1 parent 97f4016
commit 3962fd9
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -1,4 +1,5 @@
-# pytorchse3
+pytorchse3
+================
 
 <!-- WARNING: THIS FILE WAS AUTOGENERATED! DO NOT EDIT! -->
 

diff --git a/notebooks/01_se3.ipynb b/notebooks/01_se3.ipynb
@@ -109,7 +109,7 @@
     "    B = taylor_B(theta)\n",
     "    D = (1 - A / (2 * B)) / theta.pow(2)\n",
     "    \n",
-    "    V_inv = torch.eye(3) - 0.5 * log_R + D * log_R_2\n",
+    "    V_inv = torch.eye(3, dtype=A.dtype, device=A.device) - 0.5 * log_R + D * log_R_2\n",
     "    log_t_vee = torch.einsum(\"bij, bj -> bi\", V_inv, t)\n",
     "\n",
     "    return torch.concat([log_R_vee, log_t_vee], dim=-1)"

diff --git a/pytorchse3/se3.py b/pytorchse3/se3.py
@@ -58,7 +58,7 @@ def se3_log_map(T: torch.Tensor):
     B = taylor_B(theta)
     D = (1 - A / (2 * B)) / theta.pow(2)
 
-    V_inv = torch.eye(3) - 0.5 * log_R + D * log_R_2
+    V_inv = torch.eye(3, dtype=A.dtype, device=A.device) - 0.5 * log_R + D * log_R_2
     log_t_vee = torch.einsum("bij, bj -> bi", V_inv, t)
 
     return torch.concat([log_R_vee, log_t_vee], dim=-1)