diff --git a/notebooks/dynamicrafter-animating-images/dynamicrafter-animating-images.ipynb b/notebooks/dynamicrafter-animating-images/dynamicrafter-animating-images.ipynb
index 3af46e4b15e..c55034dcc6c 100644
--- a/notebooks/dynamicrafter-animating-images/dynamicrafter-animating-images.ipynb
+++ b/notebooks/dynamicrafter-animating-images/dynamicrafter-animating-images.ipynb
@@ -559,8 +559,9 @@
     "import open_clip\n",
     "\n",
     "\n",
-    "class CondStageModelWrapper:\n",
+    "class CondStageModelWrapper(torch.nn.Module):\n",
     "    def __init__(self, cond_stage_model):\n",
+    "        super().__init__()\n",
     "        self.cond_stage_model = cond_stage_model\n",
     "\n",
     "    def encode(self, tokens):\n",
@@ -571,8 +572,9 @@
     "        return torch.from_numpy(outs)\n",
     "\n",
     "\n",
-    "class EncoderFirstStageModelWrapper:\n",
+    "class EncoderFirstStageModelWrapper(torch.nn.Module):\n",
     "    def __init__(self, encode_first_stage):\n",
+    "        super().__init__()\n",
     "        self.encode_first_stage = encode_first_stage\n",
     "\n",
     "    def forward(self, x):\n",
@@ -584,8 +586,9 @@
     "        return self.forward(*args, **kwargs)\n",
     "\n",
     "\n",
-    "class EmbedderWrapper:\n",
+    "class EmbedderWrapper(torch.nn.Module):\n",
     "    def __init__(self, embedder):\n",
+    "        super().__init__()\n",
     "        self.embedder = embedder\n",
     "\n",
     "    def forward(self, x):\n",
@@ -597,8 +600,9 @@
     "        return self.forward(*args, **kwargs)\n",
     "\n",
     "\n",
-    "class CModelWrapper:\n",
+    "class CModelWrapper(torch.nn.Module):\n",
     "    def __init__(self, diffusion_model, out_channels):\n",
+    "        super().__init__()\n",
     "        self.diffusion_model = diffusion_model\n",
     "        self.out_channels = out_channels\n",
     "\n",
@@ -617,8 +621,9 @@
     "        return self.forward(*args, **kwargs)\n",
     "\n",
     "\n",
-    "class DecoderFirstStageModelWrapper:\n",
+    "class DecoderFirstStageModelWrapper(torch.nn.Module):\n",
     "    def __init__(self, decoder_first_stage):\n",
+    "        super().__init__()\n",
     "        self.decoder_first_stage = decoder_first_stage\n",
     "\n",
     "    def forward(self, x):\n",
@@ -1875,7 +1880,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.11.4"
   },
   "openvino_notebooks": {
    "imageUrl": "https://github.com/Doubiiu/DynamiCrafter/blob/main/assets/showcase/guitar0.gif?raw=true",