riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-08-01 02:31:11 +06:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

guard torch version for uint16 (#36520)

Browse Source 
* u16

* style

* fix
This commit is contained in:
Marc Sun 2025-03-05 11:27:01 +01:00 committed by GitHub
parent 66f29aaaf5
commit 752ef3fd4e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
src/transformers/modeling_utils.py
View File

@ -522,17 +522,19 @@ str_to_torch_dtype = {
"U8": torch.uint8, "U8": torch.uint8,
"I8": torch.int8, "I8": torch.int8,
"I16": torch.int16, "I16": torch.int16,
"U16": torch.uint16,
"F16": torch.float16, "F16": torch.float16,
"BF16": torch.bfloat16, "BF16": torch.bfloat16,
"I32": torch.int32, "I32": torch.int32,
"U32": torch.uint32,
"F32": torch.float32, "F32": torch.float32,
"F64": torch.float64, "F64": torch.float64,
"I64": torch.int64, "I64": torch.int64,
"U64": torch.uint64,
} }
if is_torch_greater_or_equal("2.3.0"):
str_to_torch_dtype["U16"] = torch.uint16
str_to_torch_dtype["U32"] = torch.uint32
str_to_torch_dtype["U64"] = torch.uint64
def load_state_dict( def load_state_dict(
checkpoint_file: Union[str, os.PathLike], checkpoint_file: Union[str, os.PathLike],