diff --git a/exir/serde/export_serialize.py b/exir/serde/export_serialize.py index 8ec864e489b..570af19baf8 100644 --- a/exir/serde/export_serialize.py +++ b/exir/serde/export_serialize.py @@ -148,6 +148,7 @@ def _reverse_map(d: Dict[Any, Enum]): torch.complex128: ScalarType.COMPLEXDOUBLE, torch.bool: ScalarType.BOOL, torch.bfloat16: ScalarType.BFLOAT16, + torch.uint16: ScalarType.UINT16 } diff --git a/exir/serde/schema.py b/exir/serde/schema.py index 9bbdd5b1c02..16ac724917c 100644 --- a/exir/serde/schema.py +++ b/exir/serde/schema.py @@ -35,7 +35,7 @@ class ScalarType(IntEnum): COMPLEXDOUBLE = 11 BOOL = 12 BFLOAT16 = 13 - + UINT16 = 14 class Layout(IntEnum): Unknown = 0