import torch DTYPE_MAPPING = { "float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16, "float64": torch.float64, "uint8": torch.uint8, "int8": torch.int8, "int16": torch.int16, "int32": torch.int32, "int64": torch.int64, "bool": torch.bool }