diff options
| author |  Peter Hawkins <phawkins@google.com> | 2018-05-15 10:33:48 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-05-15 10:36:11 -0700 |
| commit | a7a3bb3df12c632b81bf1b23f8405f92a0c903c3 (patch) | |
| tree | 5291bb0bed562ce40c5a13014c6c955c71e4e10a /tensorflow/compiler/jit/xla_launch_util.cc | |
| parent | 2c83cddab0cd0de78f863e47d81b4427d6519eb7 (diff) | |
Automated g4 rollback of changelist 196683444
PiperOrigin-RevId: 196691101
Diffstat (limited to 'tensorflow/compiler/jit/xla_launch_util.cc')
| -rw-r--r-- | tensorflow/compiler/jit/xla_launch_util.cc | 5 |
1 files changed, 5 insertions, 0 deletions
diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc index d0c7a93651..6a0f557627 100644 --- a/tensorflow/compiler/jit/xla_launch_util.cc +++ b/tensorflow/compiler/jit/xla_launch_util.cc @@ -195,6 +195,11 @@ void XlaComputationLaunchContext::PopulateOutputs( OP_REQUIRES_OK( ctx, ctx->allocate_output(i, const_tensor.shape(), &output_tensor)); + if (XlaTensor* xla_tensor = XlaTensor::FromTensor(output_tensor)) { + OP_REQUIRES_OK(ctx, xla_tensor->AllocateShapedBuffer( + const_tensor.dtype(), const_tensor.shape(), + client_, stream->parent()->device_ordinal())); + } Device* device = dynamic_cast<Device*>(ctx->device()); OP_REQUIRES(ctx, device != nullptr, |