add missing F8_E4M3 and F8_E5M2 keys from str_to_torch_dtype

Author: fxmarty-amd

src/transformers/modeling_utils.py

@@ -535,6 +535,10 @@ def load_sharded_checkpoint(model, folder, strict=True, prefer_safe=True):
     str_to_torch_dtype["U32"] = torch.uint32
     str_to_torch_dtype["U64"] = torch.uint64
 
+if is_torch_greater_or_equal("2.1.0"):
+    str_to_torch_dtype["F8_E4M3"] = torch.float8_e4m3fn
+    str_to_torch_dtype["F8_E5M2"] = torch.float8_e5m2
+
 
 def load_state_dict(
     checkpoint_file: Union[str, os.PathLike],