Hi, Thanks for sharing the code. I wonder how did you do the prediction in test time on the classification task? Thanks