Support Int32 datatype

tenstorrent · Mar 1, 2025 · ace4ece · ace4ece
1 parent 5fbaf8d
commit ace4ece
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tt_torch/csrc/bindings.cpp b/tt_torch/csrc/bindings.cpp
@@ -22,8 +22,10 @@ static tt::target::DataType torch_scalar_type_to_dt(torch::ScalarType st) {
     return tt::target::DataType::UInt8;
   case torch::ScalarType::Short:
     return tt::target::DataType::UInt16;
-  case torch::ScalarType::Int:
+  case torch::ScalarType::UInt32:
     return tt::target::DataType::UInt32;
+  case torch::ScalarType::Int:
+    return tt::target::DataType::Int32;
   case torch::ScalarType::Long:
     return tt::target::DataType::UInt32;
   case torch::ScalarType::Half:
@@ -54,6 +56,8 @@ static torch::ScalarType dt_to_torch_scalar_type(tt::target::DataType df) {
   case tt::target::DataType::UInt16:
     return torch::ScalarType::Short;
   case tt::target::DataType::UInt32:
+    return torch::ScalarType::UInt32;
+  case tt::target::DataType::Int32:
     return torch::ScalarType::Int;
   case tt::target::DataType::Float16:
     return torch::ScalarType::Half;