Allow training with `precision != 32`

We should remove this condition since we only support torch=1.13