diff --git a/gen_imgs.py b/gen_imgs.py
index 4e9d3ca0..fd14a3e9 100644
--- a/gen_imgs.py
+++ b/gen_imgs.py
@@ -1,5 +1,8 @@
 """Generate some plots for the pyGAM repo."""
 
+import matplotlib
+
+matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.font_manager import FontProperties
@@ -323,7 +326,6 @@ def chicago_tensor():
 
 
 def expectiles():
-
     """Generate expectiles visualization."""
     X, y = mcycle(return_X_y=True)
 
diff --git a/pygam/tests/test_GAM_methods.py b/pygam/tests/test_GAM_methods.py
index ee29de14..035611aa 100644
--- a/pygam/tests/test_GAM_methods.py
+++ b/pygam/tests/test_GAM_methods.py
@@ -174,10 +174,8 @@ def test_summary_returns_12_lines(mcycle_gam):
              known smoothing parameters, but when smoothing parameters have been estimated, the p-values
              are typically lower than they should be, meaning that the tests reject the null too readily.
     """  # noqa: E501
-    if sys.version_info.major == 2:
-        from StringIO import StringIO
-    if sys.version_info.major == 3:
-        from io import StringIO  # noqa: F811
+    from io import StringIO
+
     stdout = sys.stdout  # keep a handle on the real standard output
     sys.stdout = StringIO()  # Choose a file-like object to write to
     mcycle_gam.summary()