Add support for BFLOAT16

mlverse · Aug 30, 2023 · 421bbb8 · 421bbb8
1 parent c75eabf
commit 421bbb8
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/R/safetensors.R b/R/safetensors.R
@@ -158,7 +158,9 @@ torch_dtype_from_safe <- function(x) {
     "I8" = "int8",
     "I16" = "int16",
     "I32" = "int32",
-    "I64" = "int64"
+    "I64" = "int64",
+    "BF16" = "bfloat16",
+    cli::cli_abort("Unsupported dtype {.val {x}}")
   )
 }