fix tensor dimensions

Signed-off-by: Jan Bielak <jbielak@nvidia.com; git config --global format.signoff true>
janekb04 · Jul 31, 2023 · c8aff5f · c8aff5f
1 parent 65bd1cd
commit c8aff5f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/transformer_engine/pytorch/sequential_new/common_back/ops.py b/transformer_engine/pytorch/sequential_new/common_back/ops.py
@@ -369,8 +369,8 @@ def describe_supplementary_tensors_training(
     ) -> dict[str, TensorDescriptor]:
         return {
             "act": TensorDescriptor(self.input_shape, None, self.output_type),
-            "mu": TensorDescriptor((self.features,), None, DType.FP32),
-            "rsigma": TensorDescriptor((self.features,), None, DType.FP32),
+            "mu": TensorDescriptor((self.input_shape[-2],), None, DType.FP32),
+            "rsigma": TensorDescriptor((self.input_shape[-2],), None, DType.FP32),
         }
 
     def describe_supplementary_tensors_inference(