Flux training GPU vs CPU different results

@darsnack I have a similar problem, but I believe is not related to GPU vs CPU. The problem appears when I wrap the training code inside a function() which acctually prevents me to upgrade the code to a more dynamic custom training loop. when I do the latter, it just gives me the same result allways, but when I run it from the REPL, it runs OK. I have made a different post for this issue. I don’t know if it would be better to just continue this one… let me know @luboshanus if that’s Ok.