Skip to content

Commit e97c010

Browse files
wangshangsamtqchen
authored andcommitted
Fixing package path in tflite test (#3427)
1 parent 25bad44 commit e97c010

File tree

1 file changed

+5
-2
lines changed

1 file changed

+5
-2
lines changed

tests/python/frontend/tflite/test_forward.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,10 @@
3232
from tensorflow.python.ops import nn_ops
3333
from tensorflow.python.ops import array_ops
3434
from tensorflow.python.ops import variables
35-
from tensorflow import lite as interpreter_wrapper
35+
try:
36+
from tensorflow import lite as interpreter_wrapper
37+
except ImportError:
38+
from tensorflow.contrib import lite as interpreter_wrapper
3639

3740
import tvm.relay.testing.tf as tf_testing
3841

@@ -131,7 +134,7 @@ def compare_tflite_with_tvm(in_data, in_name, input_tensors,
131134
if init_global_variables:
132135
sess.run(variables.global_variables_initializer())
133136
# convert to tflite model
134-
converter = tf.contrib.lite.TFLiteConverter.from_session(
137+
converter = interpreter_wrapper.TFLiteConverter.from_session(
135138
sess, input_tensors, output_tensors)
136139
tflite_model_buffer = converter.convert()
137140
tflite_output = run_tflite_graph(tflite_model_buffer, in_data)

0 commit comments

Comments
 (0)