Skip to content

Commit 1187b20

Browse files
committed
Do not unconditionally pad to 4D input
1 parent b48fc58 commit 1187b20

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

core/execution/register_trt_op.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ std::vector<at::Tensor> RunCudaEngine(nvinfer1::IExecutionContext* ctx, std::pai
1515
std::vector<at::Tensor> contig_inputs{};
1616
contig_inputs.reserve(inputs.size());
1717
for (size_t i = 0; i < inputs.size(); i++) {
18-
auto dims = core::util::toDimsPad(inputs[i].sizes(), 4);
18+
auto dims = core::util::toDimsPad(inputs[i].sizes(), 1);
1919
auto shape = core::util::toVec(dims);
2020
contig_inputs.push_back(inputs[i].to(at::kCUDA).view(shape).contiguous());
2121
LOG_DEBUG("In shape:" << shape);

0 commit comments

Comments
 (0)