tests: Update fp16 test for new API

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · May 8, 2022 · 828d120 · 828d120
1 parent c275dd0
commit 828d120
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/tests/accuracy/test_fp16_accuracy.cpp b/tests/accuracy/test_fp16_accuracy.cpp
@@ -25,8 +25,10 @@ TEST_P(AccuracyTests, FP16AccuracyIsClose) {
   }
   torch::Tensor jit_accuracy = (jit_correct / jit_total) * 100;
 
-  std::vector<std::vector<int64_t>> input_shape = {{32, 3, 32, 32}};
-  auto compile_spec = torch_tensorrt::ts::CompileSpec({input_shape});
+  std::vector<int64_t> input_shape = {32, 3, 32, 32};
+  auto input = torch_tensorrt::Input(input_shape);
+  input.dtype = torch::kF16;
+  auto compile_spec = torch_tensorrt::ts::CompileSpec({input});
   compile_spec.enabled_precisions.insert(torch::kF16);
 
   auto trt_mod = torch_tensorrt::ts::compile(mod, compile_spec);