IR for the embedding bag only is below
#map0 = affine_map<(d0, d1, d2) -> (d1)>
#map1 = affine_map<(d0, d1, d2) -> (d0)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d2)>
module attributes {torch.debug_module_name = "ToyEmbeddingBag"} {
func.func @forward(%arg0: tensor<4xi64>, %arg1: tensor<2xi64>) -> tensor<2x3xf32> {
%cst = arith.constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0], [16.0, 17.0, 18.0], [19.0, 20.0, 21.0], [22.0, 23.0, 24.0], [25.0, 26.0, 27.0], [28.0, 29.0, 30.0]]> : tensor<10x3xf32>
%cst_0 = arith.constant 0.000000e+00 : f32