diff --git a/README.md b/README.md
index 3d8826a..031228c 100644
--- a/README.md
+++ b/README.md
@@ -87,10 +87,16 @@ The dataset uses the [same format as Nerfies](https://github.com/google/nerfies#
 ## Citing
 If you find our work useful, please consider citing:
 ```BibTeX
-@article{park2021hypernerf
-  author    = {Park, Keunhong and Sinha, Utkarsh and Hedman, Peter and Barron, Jonathan T. and Bouaziz, Sofien and Goldman, Dan B and Martin-Brualla, Ricardo and Seitz, Steven M.},
-  title     = {HyperNeRF: A Higher-Dimensional Representation for Topologically Varying Neural Radiance Fields},
-  journal   = {arXiv preprint arXiv:2106.13228},
-  year      = {2021},
+@article{park2021hypernerf,
+  author = {Park, Keunhong and Sinha, Utkarsh and Hedman, Peter and Barron, Jonathan T. and Bouaziz, Sofien and Goldman, Dan B and Martin-Brualla, Ricardo and Seitz, Steven M.},
+  title = {HyperNeRF: A Higher-Dimensional Representation for Topologically Varying Neural Radiance Fields},
+  journal = {ACM Trans. Graph.},
+  issue_date = {December 2021},
+  publisher = {ACM},
+  volume = {40},
+  number = {6},
+  month = {dec},
+  year = {2021},
+  articleno = {238},
 }
 ```
diff --git a/hypernerf/configs.py b/hypernerf/configs.py
index 0a70423..809ff52 100644
--- a/hypernerf/configs.py
+++ b/hypernerf/configs.py
@@ -16,7 +16,7 @@
 import dataclasses
 from typing import Any, Callable, Optional
 
-from flax import nn
+from flax import linen as nn
 import gin
 import immutabledict
 import jax.numpy
diff --git a/hypernerf/utils.py b/hypernerf/utils.py
index 5060a60..6e7e97f 100644
--- a/hypernerf/utils.py
+++ b/hypernerf/utils.py
@@ -285,7 +285,7 @@ def points_bounding_size(points):
 def shard(xs, device_count=None):
   """Split data into shards for multiple devices along the first dimension."""
   if device_count is None:
-    jax.local_device_count()
+    device_count = jax.local_device_count()
   return jax.tree_map(lambda x: x.reshape((device_count, -1) + x.shape[1:]), xs)
 
 
diff --git a/notebooks/HyperNeRF_Render_Video.ipynb b/notebooks/HyperNeRF_Render_Video.ipynb
index 40570ab..08dc65c 100644
--- a/notebooks/HyperNeRF_Render_Video.ipynb
+++ b/notebooks/HyperNeRF_Render_Video.ipynb
@@ -151,8 +151,7 @@
     {
       "cell_type": "code",
       "metadata": {
-        "id": "2QYJ7dyMcw2f",
-        "cellView": "form"
+        "id": "2QYJ7dyMcw2f"
       },
       "source": [
         "# @title Model and dataset configuration\n",
@@ -195,7 +194,7 @@
         "eval_config = configs.EvalConfig()\n",
         "\n",
         "display(Markdown(\n",
-        "    gin.config.markdownify_operative_config_str(gin.config_str())))"
+        "    gin.config.markdown(gin.config_str())))"
       ],
       "execution_count": null,
       "outputs": []
diff --git a/notebooks/HyperNeRF_Training.ipynb b/notebooks/HyperNeRF_Training.ipynb
index 593a7c8..f357b2f 100644
--- a/notebooks/HyperNeRF_Training.ipynb
+++ b/notebooks/HyperNeRF_Training.ipynb
@@ -6,8 +6,7 @@
       "name": "HyperNeRF Training.ipynb",
       "private_outputs": true,
       "provenance": [],
-      "collapsed_sections": [],
-      "toc_visible": true
+      "collapsed_sections": []
     },
     "kernelspec": {
       "name": "python3",
@@ -180,8 +179,7 @@
     {
       "cell_type": "code",
       "metadata": {
-        "id": "rz7wRm7YT9Ka",
-        "cellView": "form"
+        "id": "rz7wRm7YT9Ka"
       },
       "source": [
         "# @title Model and dataset configuration\n",
@@ -314,7 +312,7 @@
         "eval_config = configs.EvalConfig()\n",
         "\n",
         "display(Markdown(\n",
-        "    gin.config.markdownify_operative_config_str(gin.config_str())))"
+        "    gin.config.markdown(gin.config_str())))"
       ],
       "execution_count": null,
       "outputs": []
diff --git a/requirements.txt b/requirements.txt
index be3cbd0..9b3b128 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,12 +7,12 @@ jax==0.2.20
 jaxlib==0.1.71+cuda111
 Markdown==3.3.4
 matplotlib==3.4.3
-numpy==1.19.5
+numpy==1.21.0
 oauthlib==3.1.1
 opencv-python==4.5.3.56
 opt-einsum==3.3.0
 optax==0.0.9
-Pillow==8.3.2
+Pillow==9.0.0
 scipy==1.7.1
-tensorflow==2.6.0
+tensorflow==2.6.3
 tqdm==4.62.2