diff --git a/extension/aten_util/aten_bridge.cpp b/extension/aten_util/aten_bridge.cpp index 907d23e238..993f5e4077 100644 --- a/extension/aten_util/aten_bridge.cpp +++ b/extension/aten_util/aten_bridge.cpp @@ -68,6 +68,8 @@ torch::executor::ScalarType torchToExecuTorchScalarType(caffe2::TypeMeta type) { return torch::executor::ScalarType::Byte; case c10::ScalarType::Char: return torch::executor::ScalarType::Char; + case c10::ScalarType::Short: + return torch::executor::ScalarType::Short; case c10::ScalarType::Half: return torch::executor::ScalarType::Half; case c10::ScalarType::Int: @@ -95,6 +97,8 @@ c10::ScalarType execuTorchtoTorchScalarType(torch::executor::ScalarType type) { return c10::ScalarType::Byte; case torch::executor::ScalarType::Char: return c10::ScalarType::Char; + case torch::executor::ScalarType::Short: + return c10::ScalarType::Short; case torch::executor::ScalarType::Half: return c10::ScalarType::Half; case torch::executor::ScalarType::Int: