Add more logging. Hopefully improve loss normalization?