Memory leak in matplotlib

[onlinejudge95]

There is a memory leak in src/test.py caused by the following snippet

if configs.save_test_output:
    fig, axes = plt.subplots(nrows=batch_size, ncols=2, figsize=(10, 5))
    plt.tight_layout()
    ...
    plt.savefig(os.path.join(configs.saved_dir, 'batch_idx_{}_sample_idx_{}.jpg'.format(batch_idx, sample_idx)))

We need to call plt.close() post the savefig method so that plt can release memory