Consider FP8 as source

#91
2024-09-01 00:56:30 +02:00
parent abe5c6a6c9
commit 7f3ced6c3e
1 changed files with 3 additions and 0 deletions
@@ -109,6 +109,9 @@ def handle_tensors(args, writer, state_dict):

        if data.dtype == torch.bfloat16:
            data = data.to(torch.float32).numpy()
+        # this is so we don't break torch 2.0.X
+        elif data.dtype in [getattr(torch, "float8_e4m3fn", "_invalid"), getattr(torch, "float8_e5m2", "_invalid")]:
+            data = data.to(torch.float16).numpy()
        else:
            data = data.numpy()