float8 check.

Author: sayakpaul

src/diffusers/models/model_loading_utils.py

@@ -185,16 +185,15 @@ def load_model_dict_into_meta(
     accepts_dtype = "dtype" in set(inspect.signature(set_module_tensor_to_device).parameters.keys())
     empty_state_dict = model.state_dict()
     unexpected_keys = [param_name for param_name in state_dict if param_name not in empty_state_dict]
-    is_torch_e4m3fn_available = hasattr(torch, "float8_e4m3fn")
 
     for param_name, param in state_dict.items():
         if param_name not in empty_state_dict:
             continue
 
-        # We convert floating dtypes to the `dtype` passed except for float8_e4m3fn type. We also want to keep the buffers/params
+        # We convert floating dtypes to the `dtype` passed. We also want to keep the buffers/params
         # in int/uint/bool and not cast them.
-        is_param_float8_e4m3fn = is_torch_e4m3fn_available and param.dtype == torch.float8_e4m3fn
-        if torch.is_floating_point(param) and not is_param_float8_e4m3fn:
+        # TODO: revisit cases when param.dtype == torch.float8_e4m3fn
+        if torch.is_floating_point(param):
             if (
                 keep_in_fp32_modules is not None
                 and any(