Skip to content

Commit 3e27790

Browse files
committed
updated test to add dtype data to converter input
1 parent d6645d3 commit 3e27790

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

tests/util/run_graph_engine.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include "core/ir/ir.h"
55
#include "core/runtime/runtime.h"
66
#include "core/util/prelude.h"
7+
#include "core/util/trt_util.h"
78
#include "cuda_runtime_api.h"
89
#include "torch/csrc/jit/ir/ir.h"
910
#include "torch/csrc/jit/ir/irparser.h"
@@ -19,7 +20,7 @@ namespace util {
1920
std::vector<core::ir::Input> toInputs(std::vector<at::Tensor> ten) {
2021
std::vector<core::ir::Input> a;
2122
for (auto i : ten) {
22-
a.push_back(core::ir::Input(core::util::toVec(i.sizes())));
23+
a.push_back(core::ir::Input(core::util::toVec(i.sizes()), core::util::ScalarTypeToTRTDataType(i.scalar_type())));
2324
}
2425
return std::move(a);
2526
}

0 commit comments

Comments
 (0)