diff --git a/aten/src/ATen/native/mps/operations/Distributions.mm b/aten/src/ATen/native/mps/operations/Distributions.mm
index f827f81039a52..cbabf80b802e7 100644
--- a/aten/src/ATen/native/mps/operations/Distributions.mm
+++ b/aten/src/ATen/native/mps/operations/Distributions.mm
@@ -395,9 +395,13 @@ Tensor normal_mps(const Tensor& mean, const Tensor& std, c10::optional<Generator
           MPSGraphTensor *ones = [mpsGraph constantWithScalar:1.0f
                                                         shape:@[ns_numCategories, ns_numCategories]
                                                      dataType:prob_dtype];
+          MPSGraphTensor *zeroTensor = [mpsGraph constantWithScalar:0.0f
+                                                          dataType:MPSDataTypeInt32];
+          MPSGraphTensor *minusOneTensor = [mpsGraph constantWithScalar:-1.0f
+                                                              dataType:MPSDataTypeInt32];
           MPSGraphTensor *upperTriangle = [mpsGraph bandPartWithTensor:ones
-                                                              numLower:0
-                                                              numUpper:-1
+                                                        numLowerTensor:zeroTensor
+                                                        numUpperTensor:minusOneTensor
                                                                   name:nil];
           MPSGraphTensor *upperProbRange = [mpsGraph matrixMultiplicationWithPrimaryTensor:normalizedProbs
                                                                            secondaryTensor:upperTriangle