Text Classification (ULMFiT) SOTU predictions shape

I would like to train a model using all State of the Union addresses since 1790 in order to predict if a given sentence is more Republican or Democratic. I perform rather simple preprocessing: split into sentences, remove all non-alphanumeric characters and then split into words. I am using the pre-trained ULMFiT weights, and am getting fairly good accuracy (75%).

However, when I make predictions they are of a strange shape (640,70,16). I am expecting the resulting predictions to be of dimension (m,2). Can anyone tell me where I am going wrong? Thanks.

https://colab.research.google.com/drive/1JANHqrKxHZZFHFjZCxllukmWeldFbgjR

I figured it out. Thought it referenced the dataframe by column name, turns out its positional so I just switched the order of the columns :slight_smile: