pub fn embedding(weights: Tensor<2>, indices: Tensor<2, Int>) -> Tensor<3>
Applies the embedding module.