diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc
index 40729da3b4..9d801b6798 100644
--- a/orttraining/orttraining/core/graph/training_op_defs.cc
+++ b/orttraining/orttraining/core/graph/training_op_defs.cc
@@ -501,7 +501,42 @@ void RegisterTrainingOpSchemas() {
       .TypeConstraint(
           "T",
           {"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"},
-          "Constrain input and output types to numeric tensors.");
+          "Constrain input and output types to numeric tensors.")
+      .FunctionBody([]() {
+        auto nodes = ONNX_NAMESPACE::FunctionBodyHelper::BuildNodes(
+          {// nodes: {outputs, op, inputs, attributes}
+
+           // Get input shapes and dynamic reduction axes.
+           {{"shape_A"}, "Shape", {"A"}},
+           {{"shape_B"}, "Shape", {"B"}},
+           {{"axes_A", "axes_B"}, "BroadcastGradientArgs", {"shape_A", "shape_B"}},
+
+           // dA = reshape(reduce_sum(dY / B, axes_A), shape_A)
+           {{"dY_over_B"}, "Div", {"dY", "B"}},
+           {{"reduce_dA"}, "ReduceSumTraining", {"dY_over_B", "axes_A"},
+            {ONNX_NAMESPACE::MakeAttribute("noop_with_empty_axes", int64_t(1))}},
+           {{"dA"}, "Reshape", {"reduce_dA", "shape_A"}},
+
+           // dB = reshape(reduce_sum(dY * -A / (B * B)), axes_B), shape_B)
+           {{"B_squared"}, "Mul", {"B", "B"}},
+           {{"minus_A"}, "Neg", {"A"}},
+           {{"minus_A_over_B_squared"}, "Div", {"minus_A", "B_squared"}},
+           {{"pre_reduce_dB"}, "Mul", {"dY", "minus_A_over_B_squared"}},
+           {{"reduce_dB"}, "ReduceSumTraining", {"pre_reduce_dB", "axes_B"},
+            {ONNX_NAMESPACE::MakeAttribute("noop_with_empty_axes", int64_t(1))}},
+           {{"dB"}, "Reshape", {"reduce_dB", "shape_B"}}});
+
+           for (size_t contrib_node_index : {2, 4, 10}) {
+             nodes[contrib_node_index].set_domain(kMSDomain);
+           }
+           return nodes;
+      }())
+      .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
+        for (size_t i = 0; i < ctx.getNumOutputs(); ++i) {
+          propagateElemTypeFromTensorInputToOutput(ctx, 0, i);
+          propagateShapeFromInputToOutput(ctx, i + 1, i);
+        }
+      });
 
   //TODO: Move this to the right location. Its only here for quick experimentation.
   //TODO: Use the mutli weight / grad version.
diff --git a/orttraining/orttraining/test/gradient/gradient_ops_test.cc b/orttraining/orttraining/test/gradient/gradient_ops_test.cc
index 291cbe8920..b51ac8067b 100644
--- a/orttraining/orttraining/test/gradient/gradient_ops_test.cc
+++ b/orttraining/orttraining/test/gradient/gradient_ops_test.cc
@@ -294,12 +294,10 @@ TEST(GradientCheckerTest, DISABLED_MulGrad) {
   TestBroadcastableBinaryOpGrad("Mul");
 }
 
-#ifdef USE_CUDA
 TEST(GradientCheckerTest, DivGrad) {
   std::function<float(float)> transformer = [](float x) { return x > 0 ? x + 0.2f : x - 0.2f; };
   TestBroadcastableBinaryOpGrad("Div", &transformer);
 }
-#endif
 
 // TODO: Powgrad Test doesn't cover exponent
 TEST(GradientCheckerTest, PowGrad) {