I was use MNIST dataset and this is the result
Epoch 0 Train Loss 149: 0.14252015948295593
Epoch 0 Train Loss 150: 0.13949544727802277
Epoch 0 Train Loss 151: 0.14344342052936554
Epoch 0 Train Loss 152: 0.1380266398191452
Epoch 0 Train Loss 153: 0.08092612773180008
Epoch 0 Train Loss 154: 0.04462871327996254
Epoch 0 Train Loss 155: 0.026371218264102936
Epoch 0 Train Loss 156: 0.00028531771386042237
Epoch 0 Train Loss 157: 0.006431366316974163
Epoch 0 Train Loss 158: 0.007001318968832493
Epoch 0 Train Loss 159: -0.01587434858083725
Epoch 0 Train Loss 160: -0.0454784594476223
Epoch 0 Train Loss 161: -0.04419191554188728
Epoch 0 Train Loss 162: -0.06265364587306976
Epoch 0 Train Loss 163: -0.1070045754313469
Epoch 0 Train Loss 164: -0.10746034979820251
Epoch 0 Train Loss 165: -0.18705613911151886
Epoch 0 Train Loss 166: -0.1402345448732376
Epoch 0 Train Loss 167: -0.14482370018959045
Epoch 0 Train Loss 168: -0.21205617487430573
Epoch 0 Train Loss 169: -0.298073947429657
Epoch 0 Train Loss 170: -0.2425336241722107
Epoch 0 Train Loss 171: -0.3488599359989166
Epoch 0 Train Loss 172: -0.4062669575214386
Epoch 0 Train Loss 173: -0.4244926869869232
Epoch 0 Train Loss 174: -0.45135900378227234
Epoch 0 Train Loss 175: -0.49181121587753296
Epoch 0 Train Loss 176: -0.5699509382247925