// the shapes can change due to fit_output, then it may not be possible
// to return an array, even when the input was an array
if input_was_array:
nb_shapes = len(set([image.shape for image in result]))
if nb_shapes == 1:
result = np.array(result, input_dtype)
After Change
// the shapes can change due to fit_output, then it may not be possible
// to return an array, even when the input was an array
if input_was_array:
nb_shapes = len({image.shape for image in result})
if nb_shapes == 1:
result = np.array(result, input_dtype)
