Merge pull request #10 from run-ai/omer/bugfix-8f

Support float8 dtype
run-ai · Oct 15, 2024 · cb41e37 · cb41e37
2 parents 53c7b45 + 6d0d8d5
commit cb41e37
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/py/runai_model_streamer/requirements.dev b/py/runai_model_streamer/requirements.dev
@@ -1,4 +1,4 @@
-torch>=2.0.0,<3.0.0
+torch>=2.1.0,<3.0.0
 numpy==1.24.4
 safetensors==0.4.2
 humanize==4.9.0
diff --git a/py/runai_model_streamer/runai_model_streamer/safetensors_streamer/safetensors_pytorch.py b/py/runai_model_streamer/runai_model_streamer/safetensors_streamer/safetensors_pytorch.py
@@ -25,6 +25,8 @@
     "I8": torch.int8,
     "U8": torch.uint8,
     "BOOL": torch.bool,
+    "F8_E5M2": torch.float8_e5m2,
+    "F8_E4M3": torch.float8_e4m3fn,
 }