Update

[ghstack-poisoned]

Author: manuelcandales

backends/apple/metal/runtime/shims/utils.cpp

@@ -20,8 +20,10 @@ extern "C" {
 bool is_dtype_supported_in_et_metal(int32_t dtype) {
   switch (dtype) {
     case static_cast<int32_t>(SupportedDTypes::UINT8):
+    case static_cast<int32_t>(SupportedDTypes::INT32):
     case static_cast<int32_t>(SupportedDTypes::INT64):
     case static_cast<int32_t>(SupportedDTypes::FLOAT32):
+    case static_cast<int32_t>(SupportedDTypes::BOOL):
     case static_cast<int32_t>(SupportedDTypes::BFLOAT16):
       return true;
     default:
@@ -37,12 +39,8 @@ AOTITorchError validate_dtype(int32_t dtype) {
 
   ET_LOG(
       Error,
-      "Unsupported dtype: %d. Supported dtypes: %d (uint8), %d (int64), %d (float32), %d (bfloat16)",
-      dtype,
-      static_cast<int32_t>(SupportedDTypes::UINT8),
-      static_cast<int32_t>(SupportedDTypes::INT64),
-      static_cast<int32_t>(SupportedDTypes::FLOAT32),
-      static_cast<int32_t>(SupportedDTypes::BFLOAT16));
+      "Unsupported dtype: %d",
+      dtype);
   return Error::InvalidArgument;
 }
 

backends/apple/metal/runtime/shims/utils.h

@@ -22,12 +22,12 @@ enum class SupportedDTypes : int32_t {
   UINT8 = 0, // PyTorch's uint8 dtype code
   // INT8 = 1,     // PyTorch's int8 dtype code
   // INT16 = 2,    // PyTorch's int16 dtype code
-  // INT32 = 3,    // PyTorch's int32 dtype code
+  INT32 = 3, // PyTorch's int32 dtype code
   INT64 = 4, // PyTorch's int64 dtype code
   // FLOAT16 = 5,  // PyTorch's float16 dtype code
   FLOAT32 = 6, // PyTorch's float32 dtype code
   // FLOAT64 = 7,  // PyTorch's float64 dtype code
-  // BOOL = 11,    // PyTorch's bool dtype code
+  BOOL = 11, // PyTorch's bool dtype code
   BFLOAT16 = 15 // PyTorch's bfloat16 dtype code
 };