- Add hidden layer to
Classification
andRegression
heads to make training easy out of the box as currently using some backbones directly withClassification
orRegression
head was resulting in weird behavior. - Update
SAINTEmbedding
layer to make it serializable - Some bugs/typos fixes