diff --git a/src/tfscript/writing.py b/src/tfscript/writing.py
index c704061..7138eaf 100644
--- a/src/tfscript/writing.py
+++ b/src/tfscript/writing.py
@@ -1,4 +1,5 @@
 import os
+import sys
 from os.path import exists
 from tempfile import NamedTemporaryFile