MethodError: objects of type Float64 are not callable

If you print out typeof(loss), you will see that it is a Float64 and not a function as you would expect. Thus calling it fails because you can’t call a floating point number :wink:

More broadly, Zygote requires that all the code between the first use of the input x and the final calculation of loss needs to be inside that callback you pass to gradient. That means at least line 125-135 and possibly higher up in your file. I would highly, highly recommend you read through Basics · Flux and/or Overview · Flux to get an understanding of how the library works before continuing.