In your do block, the returned value is the value of the last expression. In this case, you actually return the output of a call to println, which is nothing (and it has no gradient).
The last line of the do block needs to be return allLoss, and it should work?
1 Like