diff --git a/docs/conf.py b/docs/conf.py
index 5dab84480a28d0cb1121e2c316c4763ebfd1fcf6..62d561e6949ea83f04e0311dd44a769e19962ebb 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -8,11 +8,16 @@
 # All configuration values have a default; values that are commented out
 # serve to show the default.
 
-import os
-import sys
 import inspect
+import os
 import shutil
+import sys
 
+doctest_global_setup = '''
+import pystencils
+import numpy as np
+import sympy
+'''
 __location__ = os.path.join(os.getcwd(), os.path.dirname(
     inspect.getfile(inspect.currentframe())))
 
@@ -222,21 +227,21 @@ htmlhelp_basename = 'pystencils_autodiff-doc'
 # -- Options for LaTeX output --------------------------------------------------
 
 latex_elements = {
-# The paper size ('letterpaper' or 'a4paper').
-# 'papersize': 'letterpaper',
+    # The paper size ('letterpaper' or 'a4paper').
+    # 'papersize': 'letterpaper',
 
-# The font size ('10pt', '11pt' or '12pt').
-# 'pointsize': '10pt',
+    # The font size ('10pt', '11pt' or '12pt').
+    # 'pointsize': '10pt',
 
-# Additional stuff for the LaTeX preamble.
-# 'preamble': '',
+    # Additional stuff for the LaTeX preamble.
+    # 'preamble': '',
 }
 
 # Grouping the document tree into LaTeX files. List of tuples
 # (source start file, target name, title, author, documentclass [howto/manual]).
 latex_documents = [
-  ('index', 'user_guide.tex', u'pystencils_autodiff Documentation',
-   u'Stephan Seitz', 'manual'),
+    ('index', 'user_guide.tex', u'pystencils_autodiff Documentation',
+     u'Stephan Seitz', 'manual'),
 ]
 
 # The name of an image file (relative to this directory) to place at the top of
diff --git a/tests/test_tfmad.py b/tests/test_tfmad.py
index 54daabae6c9447b9163d56e1f251db5a5086c736..309ec94cf5ac920e48dfc7b8cb09a6a65ee6c08a 100644
--- a/tests/test_tfmad.py
+++ b/tests/test_tfmad.py
@@ -179,12 +179,10 @@ def test_tfmad_gradient_check_torch():
     print('Forward output fields (to check order)')
     print(auto_diff.forward_input_fields)
 
-    a_tensor = torch.zeros(
-        *a.shape, dtype=torch.float64, requires_grad=True)
-    b_tensor = torch.zeros(
-        *b.shape, dtype=torch.float64, requires_grad=True)
-    function = auto_diff.create_tensorflow_op(
-        {a: a_tensor, b: b_tensor}, backend='torch')
+    a_tensor = torch.zeros(*a.shape, dtype=torch.float64, requires_grad=True)
+    b_tensor = torch.zeros(*b.shape, dtype=torch.float64, requires_grad=True)
+
+    function = auto_diff.create_tensorflow_op({a: a_tensor, b: b_tensor}, backend='torch')
 
     torch.autograd.gradcheck(function.apply, [a_tensor, b_tensor])