roc() crashes when called without axes argument

>>> roc(df_vbf, p_signal, p_bkg, 'pred_is_vbf')

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-54-0e04e48bf59c> in <module>
----> 1 roc(df_vbf, p_signal, p_bkg, 'pred_is_vbf')

~/repositories/httnn/lib/python3.5/site-packages/nnfwtbn/plot.py in roc(df, signal_process, background_process, discriminant, steps, selection, min, max, axes, weight)
    346     # Handle axes
    347     if axes is None:
--> 348         fig, axes = figure.subplots()
    349     if min is None:
    350         min = discriminant(df).min()

NameError: name 'figure' is not defined