Home
last modified time | relevance | path

Searched refs:engine_input_shapes (Results 1 – 1 of 1) sorted by relevance

/external/tensorflow/tensorflow/compiler/tf2tensorrt/kernels/
Dtrt_engine_op.cc99 std::vector<TensorShape>* engine_input_shapes);
335 std::vector<TensorShape>* engine_input_shapes) { in GetCompatibleCachedEngine() argument
340 *engine_input_shapes = actual_input_shapes; in GetCompatibleCachedEngine()
354 for (int i = 0; i < engine_input_shapes->size(); i++) { in GetCompatibleCachedEngine()
355 (*engine_input_shapes)[i].set_dim(0, smallest_batch_size); in GetCompatibleCachedEngine()
567 std::vector<TensorShape> engine_input_shapes(input_shapes); in GetEngine() local
568 for (int i = 0; i < engine_input_shapes.size(); i++) { in GetEngine()
570 engine_input_shapes[i].set_dim(0, max_batch_size); in GetEngine()
574 cache.emplace(engine_input_shapes, in GetEngine()
588 return cache.at(engine_input_shapes).get(); in GetEngine()
[all …]