From af71da0ac68ad036ae3748c53f47c06bf6ea69b7 Mon Sep 17 00:00:00 2001 From: ashari4 <70242157+ashari4@users.noreply.github.com> Date: Tue, 14 Dec 2021 13:12:37 -0800 Subject: [PATCH] Yield op supports bf16 (#10035) --- orttraining/orttraining/core/graph/training_op_defs.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 5289227ad3..a0b9f707e3 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -2955,7 +2955,7 @@ Return true if all elements are true and false otherwise. /*min_arity*/ 0) .Attr("non_differentiable_outputs", "The indices of the module outputs that doesn't have a gradient.", AttributeProto::INTS, OPTIONAL_VALUE) .Attr("full_shape_outputs", "The indices of the module outputs that must have full shape.", AttributeProto::INTS) - .TypeConstraint("T", OpSchema::all_tensor_types(), "Allow inputs and outputs to be any kind of tensor.") + .TypeConstraint("T", OpSchema::all_tensor_types_with_bfloat(), "Allow inputs and outputs to be any kind of tensor.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { auto non_differentiable_outputs = ctx.getAttribute("non_differentiable_outputs"); std::unordered_set non_differentiable_outputs_indices{};