@@ -25,6 +25,12 @@ auto interpolate_registrations = RegisterNodeConversionPatterns()
25
25
auto output_size = util::toDims (args[1 ].unwrapToIntList ());
26
26
27
27
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 );
28
34
} else {
29
35
LOG_DEBUG (" scale factor parameters not supported yet." );
30
36
}
@@ -34,7 +40,8 @@ auto interpolate_registrations = RegisterNodeConversionPatterns()
34
40
}).pattern({
35
41
" aten::upsample_nearest2d(Tensor self, int[2] output_size, float? scales_h=None, float? scales_w=None) -> (Tensor)" ,
36
42
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
37
- // std::raise(SIGINT);
43
+ // std::raise(SIGINT);
44
+
38
45
TRTORCH_ASSERT (args[0 ].IValue ()->isTensor (), " Input expected to be of type Tensor" );
39
46
40
47
auto in = args[0 ].ITensor ();
@@ -46,12 +53,11 @@ auto interpolate_registrations = RegisterNodeConversionPatterns()
46
53
47
54
TRTORCH_ASSERT ( (output_size.nbDims == 1 || output_size.nbDims == 2 ), " aten::upsample_nearest2d input Tensor and output size dimension mismatch" );
48
55
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);
54
58
59
+ resize_layer->setOutputDimensions (output_size);
60
+ resize_layer->setResizeMode (nvinfer1::ResizeMode::kNEAREST );
55
61
} else {
56
62
LOG_DEBUG (" scale factor parameters not supported yet." );
57
63
}
@@ -72,7 +78,11 @@ auto interpolate_registrations = RegisterNodeConversionPatterns()
72
78
73
79
TRTORCH_ASSERT ( (output_size.nbDims == 1 || output_size.nbDims == 3 ), " aten::upsample_nearest3d input Tensor and output size dimension mismatch" );
74
80
81
+ auto resize_layer = ctx->net ->addResize (*in);
82
+ TRTORCH_CHECK (resize_layer, " Unable to create interpolation (resizing) layer from node" << *n);
75
83
84
+ resize_layer->setOutputDimensions (output_size);
85
+ resize_layer->setResizeMode (nvinfer1::ResizeMode::kNEAREST );
76
86
} else {
77
87
LOG_DEBUG (" scale factor parameters not supported yet." );
78
88
}
0 commit comments