Mirrors_Framework/deepmind-research
1
0
Fork 0
mirror of https://github.com/google-deepmind/deepmind-research.git synced 2026-06-02 14:45:25 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Remove usages of apply_rng=True from Haiku code.

Browse Source 
This change is a no-op, `apply_rng=True` is the default value and for the last
year True has been the only accepted value for apply_rng. We intend to remove
this argument from Haiku in an upcoming change.

PiperOrigin-RevId: 413426614
This commit is contained in:
Tom Hennigan
2021-12-01 16:46:33 +00:00
committed by Diego de Las Casas
parent 601f12ba12
commit 03be2142f9
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
himo/himo_example.ipynb
+3 -3
View File
@@ -402,7 +402,7 @@
"\n", "\n",
" loss = jnp.mean((v_pred - r)**2)\n", " loss = jnp.mean((v_pred - r)**2)\n",
" return loss, (loss,)\n", " return loss, (loss,)\n",
" return hk.transform(loss_fn, apply_rng=True)\n", " return hk.transform(loss_fn)\n",
"\n", "\n",
"\n", "\n",
"def vnet_himo_loss(hidden_size, phi_size, model_loss_weight=3.0):\n", "def vnet_himo_loss(hidden_size, phi_size, model_loss_weight=3.0):\n",
@@ -420,7 +420,7 @@
" loss = v_loss + v_hind_loss + model_loss_weight * model_loss\n", " loss = v_loss + v_hind_loss + model_loss_weight * model_loss\n",
" return loss, (v_loss, v_hind_loss, model_loss)\n", " return loss, (v_loss, v_hind_loss, model_loss)\n",
"\n", "\n",
" return hk.transform(loss_fn, apply_rng=True)\n", " return hk.transform(loss_fn)\n",
"\n", "\n",
"\n", "\n",
"def vnet_model_loss(hidden_size, model_loss_weight=3.0):\n", "def vnet_model_loss(hidden_size, model_loss_weight=3.0):\n",
@@ -439,7 +439,7 @@
"\n", "\n",
" return loss, (test_v_loss, train_v_loss, model_loss)\n", " return loss, (test_v_loss, train_v_loss, model_loss)\n",
"\n", "\n",
" return hk.transform(loss_fn, apply_rng=True)" " return hk.transform(loss_fn)"
] ]
}, },
{ {