added explicit dtypes

atomicarchitects · Aug 14, 2024 · 66531b5 · 66531b5
1 parent 13eccf2
commit 66531b5
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/inference.cpp b/inference.cpp
@@ -13,10 +13,10 @@ int main(int argc, char *argv[]) {
     torch::inductor::AOTIModelContainerRunnerCuda *runner;
     runner = new torch::inductor::AOTIModelContainerRunnerCuda(model_path, 1);
     std::vector<torch::Tensor> inputs = {
-        torch::randn({32,1}, at::kCUDA),
-        torch::randn({32,3}, at::kCUDA),
-        torch::randn({2,50}, at::kCUDA),
-        torch::randn({32,1}, at::kCUDA)
+        torch::randn({32,1}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA)),
+        torch::randn({32,3}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA)),
+        torch::randn({2,50}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA)),
+        torch::randn({32}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA))
     };
     std::vector<torch::Tensor> outputs = runner->run(inputs);
     std::cout << "Result from the first inference:"<< std::endl;