tests(aten::zeros): Add a dtype test

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

Author: narendasan

tests/core/conversion/evaluators/test_aten_evaluators.cpp

@@ -54,5 +54,25 @@ TEST(Evaluators, ZerosEvaluatesCorrectly) {
   auto jit_results = trtorch::tests::util::EvaluateGraphJIT(g, {in});
   auto trt_results = trtorch::tests::util::EvaluateGraph(g->block(), {in});
 
+  ASSERT_TRUE(at::equal(jit_results[0].toTensor().to(at::kCUDA), trt_results[0].toTensor()));
+}
+
+TEST(Evaluators, ZerosDataTypeEvaluatesCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1 : Tensor):
+        %2 : int = prim::Constant[value=5]() # :0:0 (Float16)
+        %3 : None = prim::Constant() # :0:0
+        %4 : int[] = aten::size(%x.1) # <string>:7:9
+        %z.1 : Tensor = aten::zeros(%4, %2, %3, %3, %3) # experiments/test_zeros.py:8:12
+        return (%z.1))IR";
+
+  auto in = at::randint(1, 10, {1, 5, 5, 5}, {at::kCUDA});
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto jit_results = trtorch::tests::util::EvaluateGraphJIT(g, {in});
+  auto trt_results = trtorch::tests::util::EvaluateGraph(g->block(), {in});
+
   ASSERT_TRUE(at::equal(jit_results[0].toTensor().to(at::kCUDA), trt_results[0].toTensor()));
 }