saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 17:15:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Yield op supports bf16 (#10035)

Browse source
This commit is contained in:
ashari4 2021-12-14 13:12:37 -08:00 committed by GitHub
parent 703becd796
commit af71da0ac6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
orttraining/orttraining/core/graph/training_op_defs.cc
View file

@ -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<size_t> non_differentiable_outputs_indices{};