A fix for a bug found while using automated tests.
In case of single cells the plotting arrays should be of shape [trains, x, y] by indexing the cell.
Testing CALLAB collected JF runs. and running the notebook interactively.
@schmidtp @danilevc @kluyvert