fix(//tests/cpp): Fix the BERT C++ test

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

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);
       }
 

core/runtime/execute_engine.cpp

@@ -351,7 +351,6 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
     // If in CUDAGraph mode, results need to be copied to the result buffers (on caller stream)
     for (size_t o = 0; o < compiled_engine->output_buffers.size(); o++) {
       outputs[o].copy_(compiled_engine->output_buffers[o], false);
-    }
   }
 
   if (compiled_engine->profile_execution) {

tests/cpp/test_compiled_modules.cpp

@@ -5,7 +5,11 @@ TEST_P(CppAPITests, CompiledModuleIsClose) {
   std::vector<torch::jit::IValue> trt_inputs_ivalues;
   std::vector<torch_tensorrt::Input> 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]);

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)