check torch version

Author: SunMarc

src/accelerate/utils/modeling.py

@@ -43,7 +43,7 @@
 from .memory import clear_device_cache
 from .offload import load_offloaded_weight, offload_weight, save_offload_index
 from .tqdm import is_tqdm_available, tqdm
-from .versions import compare_versions
+from .versions import compare_versions, is_torch_version
 
 
 if is_npu_available(check_device=False):
@@ -161,8 +161,10 @@ def dtype_byte_size(dtype: torch.dtype):
         return 1 / 4
     elif dtype == CustomDtype.INT4:
         return 1 / 2
-    elif dtype in [CustomDtype.FP8, torch.float8_e4m3fn]:
+    elif dtype == CustomDtype.FP8:
         return 1
+    elif is_torch_version(">=", "2.1.0") and dtype == torch.float8_e4m3fn:
+        return 1 
     bit_search = re.search(r"[^\d](\d+)$", str(dtype))
     if bit_search is None:
         raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")