diff --git a/check/nbformat b/check/nbformat
index 2a123c9ffb8..1496a8e0a6d 100755
--- a/check/nbformat
+++ b/check/nbformat
@@ -15,7 +15,7 @@
 ################################################################################
 
 only_print=1
-for arg in $@; do
+for arg in "$@"; do
     if [[ "${arg}" == "--apply" ]]; then
         only_print=0
     else
@@ -25,29 +25,33 @@ for arg in $@; do
 done
 
 # Get the working directory to the repo root.
-cd "$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
-cd "$(git rev-parse --show-toplevel)"
+cd "$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" || exit
+cd "$(git rev-parse --show-toplevel)" || exit
 
 # Install a pinned version of tf-docs
 TF_DOCS_COMMIT=b3dc8a922d8bdc6e998a8ad4f4953359dc6e576a
-TF_DOCS_INSTALLED=$(pip list | grep  -c $TF_DOCS_COMMIT)
+TF_DOCS_INSTALLED=$(pip list | grep -c "$TF_DOCS_COMMIT")
 
 if [ "$TF_DOCS_INSTALLED" != "1" ]; then
   pip install -U git+https://github.com/tensorflow/docs@$TF_DOCS_COMMIT
 fi
 
 FORMAT_CMD="python3 -m tensorflow_docs.tools.nbfmt --indent=1"
+TFDOCS_FORMAT="black-nb"
 
 # Test the notebooks
-unformatted=$($FORMAT_CMD --test docs 2>&1 | grep "\- docs" || true)
+black_unformatted="$("$TFDOCS_FORMAT" --check docs 2>&1 | grep "would be reformatted")"
+unformatted="$("$FORMAT_CMD" --test docs 2>&1 | grep "\- docs" || true)"
 needed_changes=0
-if [ ! -z "${unformatted}" ]; then
+if [ -n "${unformatted}" ] || [ -n "${black_unformatted}" ]; then
     needed_changes=1
     if (( only_print == 0 )); then
-        $FORMAT_CMD docs
+        "$TFDOCS_FORMAT" docs
+        "$FORMAT_CMD" docs
     else
         echo -e "\033[31mThe following notebooks require formatting\033[0m."
         echo "${unformatted}"
+        echo "${black_unformatted}"
     fi
 fi
 
@@ -58,4 +62,4 @@ elif (( only_print == 1 )); then
     exit 1
 else
     echo -e "\033[33mReformatted changed notebooks\033[0m."
-fi
\ No newline at end of file
+fi
diff --git a/dev_tools/conf/pip-list-dev-tools.txt b/dev_tools/conf/pip-list-dev-tools.txt
index e5eafb200f9..9dabb80f989 100644
--- a/dev_tools/conf/pip-list-dev-tools.txt
+++ b/dev_tools/conf/pip-list-dev-tools.txt
@@ -3,6 +3,7 @@
 asv
 virtualenv
 black==20.8b1
+black-nb~=0.4.0
 mypy~=0.782.0
 pylint~=2.6.0
 pytest~=6.2.2
@@ -39,4 +40,4 @@ rstcheck~=3.3.1
 freezegun~=0.3.15
 
 # for python 3.7 and below needs to be installed
-importlib-metadata; python_version < '3.8'
\ No newline at end of file
+importlib-metadata; python_version < '3.8'