Minor fix in JAX_porting_PyTorch_model tutorial

jax-ml · Nov 8, 2024 · 369bed6 · 369bed6
1 parent 3cc817a
commit 369bed6
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/docs/JAX_porting_PyTorch_model.ipynb b/docs/JAX_porting_PyTorch_model.ipynb
@@ -592,9 +592,12 @@
     "        self.p = p\n",
     "        self.mode = mode\n",
     "        self.deterministic = False\n",
+    "        self.rngs = rngs\n",
     "\n",
     "    def __call__(self, x: jax.Array) -> jax.Array:\n",
-    "        return stochastic_depth(x, self.p, self.mode, self.deterministic)"
+    "        return stochastic_depth(\n",
+    "            x, self.p, self.mode, self.deterministic, rngs=self.rngs\n",
+    "        )"
    ]
   },
   {

diff --git a/docs/JAX_porting_PyTorch_model.md b/docs/JAX_porting_PyTorch_model.md
@@ -361,9 +361,12 @@ class StochasticDepth(nnx.Module):
         self.p = p
         self.mode = mode
         self.deterministic = False
+        self.rngs = rngs
 
     def __call__(self, x: jax.Array) -> jax.Array:
-        return stochastic_depth(x, self.p, self.mode, self.deterministic)
+        return stochastic_depth(
+            x, self.p, self.mode, self.deterministic, rngs=self.rngs
+        )
 ```
 
 ```{code-cell} ipython3