diff --git a/examples/square.py b/examples/square.py
index 64b47c49b679f6d1c467022bed18444dbb9b9a3f..14e0bd272b49972eb3e27556e2a819b479ea53c8 100644
--- a/examples/square.py
+++ b/examples/square.py
@@ -4,6 +4,7 @@ kwant.Builder.
 
 from __future__ import division
 import numpy as np
+from matplotlib import pyplot
 import kwant
 from kwant.physics.selfenergy import square_self_energy
 
@@ -105,11 +106,10 @@ def main():
     data = [kwant.solve(sys, energy).transmission(1, 0)
             for energy in energies]
 
-    import pylab
-    pylab.plot(energies, data)
-    pylab.xlabel("energy [in units of t]")
-    pylab.ylabel("conductance [in units of e^2/h]")
-    pylab.show()
+    pyplot.plot(energies, data)
+    pyplot.xlabel("energy [in units of t]")
+    pyplot.ylabel("conductance [in units of e^2/h]")
+    pyplot.show()
 
 
 if __name__ == '__main__':