diff --git a/environment.yml b/environment.yml
index cf41740..f597e1b 100644
--- a/environment.yml
+++ b/environment.yml
@@ -7,6 +7,8 @@ dependencies:
   - python>=3.6
   - pytorch>=1.6
   - torchvision
+  - execnb
+  - nbformat
   - pip
   - pip:
     - -r requirements.txt
diff --git a/tools/clean.py b/tools/clean.py
index 298e6b2..8417184 100755
--- a/tools/clean.py
+++ b/tools/clean.py
@@ -4,6 +4,7 @@ import nbformat
 from nbdev.export import *
 from nbdev.clean import *
 from fastcore.all import *
+from execnb.nbio import *
 
 _re_header = re.compile(r'^#+\s+\S+')
 _re_clean  = re.compile(r'^\s*#\s*clean\s*')