WINO=1 python examples/beautiful_mnist.py has lower test_accuracy

[chenyuxyz]

I can hit RuntimeError: Error Domain=MTLCommandBufferErrorDomain Code=1 "Discarded (victim of GPU error/recovery) on M1 Max pretty consistently with WINO=1 PYTHONPATH=. python examples/beautiful_mnist.py. Seems fine with JIT=2.
Also with WINO=1 when it runs it only trains to 95%

[geohot]

I don't see the crash on M3, but I do see the low accuracy.

[chenyuxyz]

weird, rebooted and it does not crash. gpu was likely in a bad state.

update the issue for low accuracy

[michbogos]

I've also seen lower accuracy with IMAGE=1.

[chenyuxyz]

WINO and IMAGE both are affected by LAZYCACHE with JIT, and eval step is not properly captured by jit.

As a workaround either commenting out TinyJit of get_test_acc, or adding LAZYCACHE=0 works.