feature_type};
      auto training_op = rewriter.create<BatchNormGradOp>(
          loc, operand_types, act, scale, mean, var, grad, op.getEpsilon(),
          feature_dim);

      x_backprop = training_op.getResult(0);

      scale_backprop = training_op.getResult(1);

      offset_backprop = training_op.getResult(2);
    } else {  // inference