Searched refs:device_ordinal_tensor (Results 1 – 3 of 3) sorted by relevance
78 const Tensor& device_ordinal_tensor = ctx->input(1); in ComputeAsync() local80 ctx, TensorShapeUtils::IsScalar(device_ordinal_tensor.shape()), in ComputeAsync()82 device_ordinal_tensor.shape().DebugString()), in ComputeAsync()84 const int device_ordinal = device_ordinal_tensor.flat<int64_t>()(0); in ComputeAsync()235 const Tensor& device_ordinal_tensor = ctx->input(ctx->num_inputs() - 1); in Compute() local237 ctx, TensorShapeUtils::IsScalar(device_ordinal_tensor.shape()), in Compute()239 device_ordinal_tensor.shape().DebugString())); in Compute()240 const int device_ordinal = device_ordinal_tensor.flat<int64_t>()(0); in Compute()
125 const Tensor* device_ordinal_tensor; in Compute() local126 OP_REQUIRES_OK(ctx, ctx->input("device_ordinal", &device_ordinal_tensor)); in Compute()127 device_ordinal_ = device_ordinal_tensor->flat<int32>()(0); in Compute()
122 const Tensor& device_ordinal_tensor = ctx->input(0); in RunTransfer() local123 const int device_ordinal = device_ordinal_tensor.scalar<int32>()(); in RunTransfer()