support ascend rms_norm

DeepLink-org · Jan 18, 2024 · b1cced5 · b1cced5
1 parent a39abfd
commit b1cced5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ext_op/example_ext.cpp b/ext_op/example_ext.cpp
@@ -44,8 +44,8 @@ auto extRmsNorm(const at::Tensor& input,
   at::OptionalIntArrayRef normalized_shape_at =
       optionalIntArrayToIntArrayRefOrDefault(normalized_shape, weight.sizes());
   auto input_shape = input.sizes();
-  std::vector<int64_t> input_size(input_shape.size(), 1);
-  std::copy(input_shape.begin(), input_shape.end() - 1, input_size.begin());
+  std::vector<int64_t> input_size(input_shape.begin(), input_shape.end());
+  input_size.back() = 1;
   auto inv_rms = at::empty(input_size, input.options());
   auto output = at::empty_like(input);
   callDiopi(diopiRMSNorm, output, inv_rms, input, normalized_shape_at, weight,