Unsqueeze single pose matrices

Author: eigenvivek

diffdrr/pose.py

@@ -19,6 +19,8 @@ class RigidTransform(torch.nn.Module):
 
     def __init__(self, matrix):
         super().__init__()
+        if specimen.extrinsic.matrix.dim() == 2:
+            matrix = matrix.unsqueeze(0)
         self.register_buffer("matrix", matrix)
 
     def __len__(self):

notebooks/api/06_pose.ipynb

@@ -95,6 +95,8 @@
     "\n",
     "    def __init__(self, matrix):\n",
     "        super().__init__()\n",
+    "        if specimen.extrinsic.matrix.dim() == 2:\n",
+    "            matrix = matrix.unsqueeze(0)\n",
     "        self.register_buffer(\"matrix\", matrix)\n",
     "\n",
     "    def __len__(self):\n",