Skip to content

Commit 4da65b3

Browse files
authored
Merge pull request #436 from guoruoqian/replication_padXd_fix_bug
fix bug of replication_pad when input is dynamic
2 parents 7be368f + 173918f commit 4da65b3

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

core/conversion/converters/impl/replication_pad.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,9 @@ bool replication_padXd(ConversionCtx* ctx, const torch::jit::Node* n, args& args
6767
at::Tensor dimValue = torch::tensor({axis}, torch::kInt32);
6868
auto dimTensor = tensor_to_const(ctx, dimValue);
6969
indicesTensor = ctx->net->addGather(*shapeTensor, *dimTensor, 0)->getOutput(0);
70+
auto oneTensor = tensor_to_const(ctx, torch::tensor({1}, torch::kInt32));
71+
indicesTensor =
72+
ctx->net->addElementWise(*indicesTensor, *oneTensor, nvinfer1::ElementWiseOperation::kSUB)->getOutput(0);
7073
} else {
7174
auto indices = torch::tensor({inDims.d[axis] - 1}, torch::kInt32);
7275
indicesTensor = tensor_to_const(ctx, indices);

0 commit comments

Comments
 (0)