diff --git a/diffpose/calibration.py b/diffpose/calibration.py
index d3ca073..2d334f9 100644
--- a/diffpose/calibration.py
+++ b/diffpose/calibration.py
@@ -17,6 +17,8 @@
 
 @beartype
 class RigidTransform(Transform3d):
+    """Wrapper of pytorch3d.transforms.Transform3d with extra functionalities."""
+
     @jaxtyped
     def __init__(
         self,
@@ -27,7 +29,6 @@ def __init__(
         device=None,
         dtype=torch.float32,
     ):
-        """Wrapper of pytorch3d.transforms.Transform3d with extra functionalities."""
         if device is None and (R.device == t.device):
             device = R.device
 
diff --git a/notebooks/api/02_calibration.ipynb b/notebooks/api/02_calibration.ipynb
index 148ae71..5045092 100644
--- a/notebooks/api/02_calibration.ipynb
+++ b/notebooks/api/02_calibration.ipynb
@@ -104,6 +104,8 @@
     "\n",
     "@beartype\n",
     "class RigidTransform(Transform3d):\n",
+    "    \"\"\"Wrapper of pytorch3d.transforms.Transform3d with extra functionalities.\"\"\"\n",
+    "\n",
     "    @jaxtyped\n",
     "    def __init__(\n",
     "        self,\n",
@@ -114,7 +116,6 @@
     "        device=None,\n",
     "        dtype=torch.float32,\n",
     "    ):\n",
-    "        \"\"\"Wrapper of pytorch3d.transforms.Transform3d with extra functionalities.\"\"\"\n",
     "        if device is None and (R.device == t.device):\n",
     "            device = R.device\n",
     "\n",