<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def get_which_data_ycols(model, which_data_ycols):
 """ Helper to get the data columns to plot. """

    if which_data_ycols == 'all' or which_data_ycols is None:
        return np.arange(model.output_dim)
    return which_data_ycols