未验证 提交 c042d8f2 编写于 作者: W wenbin 提交者: GitHub

fix ernie (#37839)

* fix

* Update stack_op_plugin.cu

comments
上级 ead81230
...@@ -107,8 +107,13 @@ bool StackPluginDynamic::supportsFormatCombination( ...@@ -107,8 +107,13 @@ bool StackPluginDynamic::supportsFormatCombination(
const nvinfer1::PluginTensorDesc& in = in_out[pos]; const nvinfer1::PluginTensorDesc& in = in_out[pos];
if (pos == 0) { if (pos == 0) {
if (with_fp16_) { if (with_fp16_) {
return (in.type == nvinfer1::DataType::kFLOAT || return (
in.type == nvinfer1::DataType::kHALF) && // It's workaround for ernie fix len model.
// Enabling float, half on the same time will cause trt hang.
#if IS_TRT_VERSION_LT(8000)
in.type == nvinfer1::DataType::kFLOAT ||
#endif
in.type == nvinfer1::DataType::kHALF) &&
(in.format == nvinfer1::TensorFormat::kLINEAR); (in.format == nvinfer1::TensorFormat::kLINEAR);
} else { } else {
return (in.type == nvinfer1::DataType::kFLOAT) && return (in.type == nvinfer1::DataType::kFLOAT) &&
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册