Having some trouble recreating the `solve_image` result in lesson 8

I’m following along with the lesson 8 work book, and I’m having trouble recreating the result from solve_image. I think that my problem is probably somehow dependency related. I am running the examples with newer versions of several things. E.g. Python 3.6, keras 2.

The issue seems to be that the loss values are outputting matrices instead of scalars. Any idea what would cause this?

I found a way to get better output from solve_image. Since my loss was outputting a tensor, and the example code was outputting scalars, I decided that the loss function must have changed between versions. What I have is this:

def mean_squared_error(y_true, y_pred):
    return K.mean(K.square(y_pred - y_true), axis=-1)

Looking at the docs it appeared as though the axis was the reason why the loss would return the wrong thing, so I tried with this:

def f2(y_pred, y_true):
    return K.mean(K.square(y_pred - y_true), axis=None)

And then initialized the loss variable with that and got a much better solve_image result than the grey or orange image with a few specks on it: