Small error in printing status

[navid-mahmoudian]

Hello,

Thank you for this nice library. I noticed that in your CompressAI/examples/train.py example, inside the test_epoch function when you are printing 'Average losses:' you are actually printing the last values rather than average values. I think

    print(f'Test epoch {epoch}: Average losses:'
          f'\tLoss: {loss.val:.3f} |'
          f'\tMSE loss: {mse_loss.val:.3f} |'
          f'\tBpp loss: {bpp_loss.val:.2f} |'
          f'\tAux loss: {aux_loss.val:.2f}\n')

Must be changed to

    print(f'Test epoch {epoch}: Average losses:'
          f'\tLoss: {loss.avg:.3f} |'
          f'\tMSE loss: {mse_loss.avg:.3f} |'
          f'\tBpp loss: {bpp_loss.avg:.2f} |'
          f'\tAux loss: {aux_loss.avg:.2f}\n')

I know that it is very subtle, but just for the sake of completeness I am reporting this error.

[navid-mahmoudian]

Similar problem is for return value of test_epoch function, i.e. rather than return loss.valyou should return loss.avg.

[jbegaint]

Thanks for the report Navid ! We'll push a fix soon.

[fracape]

Thanks for spotting this, Navid!