ptrendx commented on a change in pull request #19011:
URL: https://github.com/apache/incubator-mxnet/pull/19011#discussion_r486614971
##########
File path: src/operator/subgraph/tensorrt/tensorrt.cc
##########
@@ -289,6 +293,21 @@ OpStatePtr TRTCreateState(const nnvm::NodeAttrs& attrs,
Context ctx,
} else if (it_inputs != inputs_to_idx.end()) {
shape_inputs[i] = in_shape[it_inputs->second];
dtype_inputs[i] = in_type[it_inputs->second];
+ if (tensorrt_int8) {
+ int dtype_size;
+ if (dtype_inputs[i] == mshadow::kFloat32) {
+ dtype_size = 4;
+ } else if (dtype_inputs[i] == mshadow::kFloat16) {
+ dtype_size = 2;
+ } else {
+ LOG(FATAL) << "TensorRT op supports only float32 and float16
inputs.";
+ }
+ size_t buffer_size = shape_inputs[i].Size() * dtype_size;
+ void *ptr;
+ cudaMalloc(&ptr, buffer_size);
Review comment:
Can we either use the storage manager or at least check for errors here?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
