The calculation of average loss doesn’t seem right to me:

avg_loss = Nx.add(Nx.mean(batch_loss), loss) |> Nx.divide(j + 1)

I believe loss is the average loss returned by the previous iteration, which includes j samples so far, so I would think we’d need to weight that by multiplying by j before dividing by j + 1… something like this?

avg_loss = Nx.add(Nx.mean(batch_loss), Nx.multiply(loss, j)) |> Nx.divide(j + 1)

or am I confusing myself?