diff options
| author |  Jared Duke <jdduke@google.com> | 2018-07-26 10:53:21 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-26 10:56:43 -0700 |
| commit | 6e658c0a5ca77677a954a34fb98f241c592c970d (patch) | |
| tree | b645103887539af5232b3f70d80a2eb9b77ed63a /tensorflow/contrib/lite/model.cc | |
| parent | 0a3155f7fbf56df5e81c7cbf35afd45173359635 (diff) | |
Add one_hot op support to TFLite
PiperOrigin-RevId: 206185190
Diffstat (limited to 'tensorflow/contrib/lite/model.cc')
| -rw-r--r-- | tensorflow/contrib/lite/model.cc | 8 |
1 files changed, 8 insertions, 0 deletions
diff --git a/tensorflow/contrib/lite/model.cc b/tensorflow/contrib/lite/model.cc index c6869feb16..5814cddc5b 100644 --- a/tensorflow/contrib/lite/model.cc +++ b/tensorflow/contrib/lite/model.cc @@ -730,6 +730,14 @@ TfLiteStatus ParseOpData(const Operator* op, BuiltinOperator op_type, *builtin_data = static_cast<void*>(params); break; } + case BuiltinOperator_ONE_HOT: { + auto* params = MallocPOD<TfLiteOneHotParams>(); + if (auto* schema_params = op->builtin_options_as_OneHotOptions()) { + params->axis = schema_params->axis(); + } + *builtin_data = static_cast<void*>(params); + break; + } // Below are the ops with no builtin_data strcture. case BuiltinOperator_BATCH_TO_SPACE_ND: |