Error Messages:

RuntimeError: Invalid DISPLAY variable

Solution:

add the following codes

plt.switch_backend('agg')

after this line

import matplotlib.pyplot as plt

Done!