pub fn load_layer_norm<B: Backend>(
    path: &str,
    device: &B::Device
) -> Result<LayerNorm<B>, Box<dyn Error>>