@kirk86 I played a bit with your script and indeed I can confirm that it is unexpected hard to get at quick convergence to near 0 loss for the single-test-batch. It immediately works for me with a significantly smaller learning rate like 1e-3 and small batch-sizes like 1,4 or 32  (no time to test others).