From ee21db0dc03db1c6b030beee4a86ec0d341017dd Mon Sep 17 00:00:00 2001 From: Naren Dasan Date: Mon, 5 Aug 2024 17:38:23 -0700 Subject: [PATCH] fix(//tests/cpp): Fix the BERT C++ test Signed-off-by: Naren Dasan Signed-off-by: Naren Dasan --- core/conversion/converters/impl/layer_norm.cpp | 2 -- tests/cpp/test_compiled_modules.cpp | 6 +++++- tests/py/ts/models/test_models.py | 4 ++-- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/core/conversion/converters/impl/layer_norm.cpp b/core/conversion/converters/impl/layer_norm.cpp index e83a56593e..4bb1c1211b 100644 --- a/core/conversion/converters/impl/layer_norm.cpp +++ b/core/conversion/converters/impl/layer_norm.cpp @@ -32,7 +32,6 @@ auto layer_norm_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() gamma = tensor_to_const(ctx, gamma_torch_tensor); } else { gamma = args[2].ITensorOrFreeze(ctx); - // gamma = broadcast(ctx, n, gamma, input_shape_vec.size(), "gamma"); gamma = add_expand(ctx, gamma, input_shape); } @@ -43,7 +42,6 @@ auto layer_norm_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() beta = tensor_to_const(ctx, beta_torch_tensor); } else { beta = args[3].ITensorOrFreeze(ctx); - // beta = broadcast(ctx, n, beta, input_shape_vec.size(), "beta"); beta = add_expand(ctx, beta, input_shape); } diff --git a/tests/cpp/test_compiled_modules.cpp b/tests/cpp/test_compiled_modules.cpp index 7def168249..d6a4378c65 100644 --- a/tests/cpp/test_compiled_modules.cpp +++ b/tests/cpp/test_compiled_modules.cpp @@ -5,7 +5,11 @@ TEST_P(CppAPITests, CompiledModuleIsClose) { std::vector trt_inputs_ivalues; std::vector shapes; for (uint64_t i = 0; i < input_shapes.size(); i++) { - auto in = at::randint(5, input_shapes[i], {at::kCUDA}).to(input_types[i]); + auto in = at::randn(input_shapes[i], {at::kCUDA}).to(input_types[i]); + if (input_types[i] == at::kInt || input_types[i] == at::kLong) { + auto in = at::randint(0, 2, input_shapes[i], {at::kCUDA}).to(input_types[i]); + } + jit_inputs_ivalues.push_back(in.clone()); trt_inputs_ivalues.push_back(in.clone()); auto in_spec = torch_tensorrt::Input(input_shapes[i]); diff --git a/tests/py/ts/models/test_models.py b/tests/py/ts/models/test_models.py index 1d5c3bae3b..1b5fd80b44 100644 --- a/tests/py/ts/models/test_models.py +++ b/tests/py/ts/models/test_models.py @@ -93,7 +93,7 @@ def test_efficientnet_b0(self): ) def test_bert_base_uncased(self): - self.model = cm.BertModule().cuda() + self.model = cm.BertModule() self.input = torch.randint(0, 2, (1, 14), dtype=torch.int32).to("cuda") compile_spec = { @@ -116,7 +116,7 @@ def test_bert_base_uncased(self): "enabled_precisions": {torch.float}, "truncate_long_and_double": True, } - with torchtrt.logging.errors(): + with torchtrt.logging.debug(): trt_mod = torchtrt.ts.compile(self.model, **compile_spec) model_outputs = self.model(self.input, self.input)