minor fix for bfloat16 (#7003)

huggingface · Jun 25, 2024 · 1e1d313 · 1e1d313
1 parent bfb0a41
commit 1e1d313
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/datasets/utils/_dill.py b/src/datasets/utils/_dill.py
@@ -165,7 +165,7 @@ def _save_torchTensor(pickler, obj):
     def create_torchTensor(np_array, dtype=None):
         tensor = torch.from_numpy(np_array)
         if dtype:
-            tensor = tensor.type(torch.bfloat16)
+            tensor = tensor.type(dtype)
         return tensor
 
     log(pickler, f"To: {obj}")