Fixing format.

huggingface · Jul 26, 2024 · e4811ac · e4811ac
1 parent 146fd4b
commit e4811ac
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/bindings/python/py_src/safetensors/torch.py b/bindings/python/py_src/safetensors/torch.py
@@ -173,7 +173,9 @@ def save_model(
         raise ValueError(msg)
 
 
-def load_model(model: torch.nn.Module, filename: Union[str, os.PathLike], strict: bool = True, device: Union[str, int] = "cpu") -> Tuple[List[str], List[str]]:
+def load_model(
+    model: torch.nn.Module, filename: Union[str, os.PathLike], strict: bool = True, device: Union[str, int] = "cpu"
+) -> Tuple[List[str], List[str]]:
     """
     Loads a given filename onto a torch model.
     This method exists specifically to avoid tensor sharing issues which are
@@ -340,6 +342,7 @@ def load(data: bytes) -> Dict[str, torch.Tensor]:
     flat = deserialize(data)
     return _view2torch(flat)
 
+
 # torch.float8 formats require 2.1; we do not support these dtypes on earlier versions
 _float8_e4m3fn = getattr(torch, "float8_e4m3fn", None)
 _float8_e5m2 = getattr(torch, "float8_e5m2", None)

diff --git a/bindings/python/tests/test_pt_comparison.py b/bindings/python/tests/test_pt_comparison.py
@@ -66,7 +66,7 @@ def test_odd_dtype(self):
 
     def test_odd_dtype_fp8(self):
         if torch.__version__ < "2.1":
-            return # torch.float8 requires 2.1
+            return  # torch.float8 requires 2.1
 
         data = {
             "test1": torch.tensor([-0.5], dtype=torch.float8_e4m3fn),