It looks like you must have changed the “return” statement in the two_layer_model function. That is part of the template code that was given to you and there should be no need to modify it. Why did you change that line? If you “copy/pasted” the entire function from an earlier version of the course (pre April 2021), then please be aware that doesn’t work: many things changed in April 2021.
There is a procedure documented on the FAQ Thread to get back to a clean copy of the notebook for comparison purposes.