Skip to content

Commit 7f12160

Browse files
committed
feat(//core/conversion/converters/impl): logic implemented
Signed-off-by: Abhiram Iyer <abhi.iyer.ai@gmail.com> Signed-off-by: Abhiram Iyer <abhirami@nvidia.com>
1 parent 1b50484 commit 7f12160

File tree

1 file changed

+16
-6
lines changed

1 file changed

+16
-6
lines changed

core/conversion/converters/impl/interpolate.cpp

+16-6
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,12 @@ auto interpolate_registrations = RegisterNodeConversionPatterns()
2525
auto output_size = util::toDims(args[1].unwrapToIntList());
2626

2727
TRTORCH_ASSERT(output_size.nbDims == 1, "aten::upsample_nearest1d input Tensor and output size dimension mismatch");
28+
29+
auto resize_layer = ctx->net->addResize(*in);
30+
TRTORCH_CHECK(resize_layer, "Unable to create interpolation (resizing) layer from node" << *n);
31+
32+
resize_layer->setOutputDimensions(output_size);
33+
resize_layer->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
2834
} else {
2935
LOG_DEBUG("scale factor parameters not supported yet.");
3036
}
@@ -34,7 +40,8 @@ auto interpolate_registrations = RegisterNodeConversionPatterns()
3440
}).pattern({
3541
"aten::upsample_nearest2d(Tensor self, int[2] output_size, float? scales_h=None, float? scales_w=None) -> (Tensor)",
3642
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
37-
// std::raise(SIGINT);
43+
// std::raise(SIGINT);
44+
3845
TRTORCH_ASSERT(args[0].IValue()->isTensor(), "Input expected to be of type Tensor");
3946

4047
auto in = args[0].ITensor();
@@ -46,12 +53,11 @@ auto interpolate_registrations = RegisterNodeConversionPatterns()
4653

4754
TRTORCH_ASSERT( (output_size.nbDims == 1 || output_size.nbDims == 2), "aten::upsample_nearest2d input Tensor and output size dimension mismatch");
4855

49-
nvinfer1::ILayer* new_layer;
50-
51-
52-
53-
//util::toDims(args[1].unwrapToIntList());
56+
auto resize_layer = ctx->net->addResize(*in);
57+
TRTORCH_CHECK(resize_layer, "Unable to create interpolation (resizing) layer from node" << *n);
5458

59+
resize_layer->setOutputDimensions(output_size);
60+
resize_layer->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
5561
} else {
5662
LOG_DEBUG("scale factor parameters not supported yet.");
5763
}
@@ -72,7 +78,11 @@ auto interpolate_registrations = RegisterNodeConversionPatterns()
7278

7379
TRTORCH_ASSERT( (output_size.nbDims == 1 || output_size.nbDims == 3), "aten::upsample_nearest3d input Tensor and output size dimension mismatch");
7480

81+
auto resize_layer = ctx->net->addResize(*in);
82+
TRTORCH_CHECK(resize_layer, "Unable to create interpolation (resizing) layer from node" << *n);
7583

84+
resize_layer->setOutputDimensions(output_size);
85+
resize_layer->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
7686
} else {
7787
LOG_DEBUG("scale factor parameters not supported yet.");
7888
}

0 commit comments

Comments
 (0)