Fixed bug in classification (I think)

Thanks Nasir