#158: Fix one-hot-encoding layer

-When a single samples was passed through the categorical head, it dropped the batch dimension. We need a squeeze before calling F.one_hot

Closes #158 (closed)