Llama3.1 405B using FP8 TypeError: couldn't find storage object Float8_e4m3fnStorage

File "/workdir/user_repository/inference/local_deploy_demo.py", line 41, in load_model self.model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", trust_remote_code=True, File "/usr/local/lib/python3.9/site-packages/transformers/models/auto/auto_factory.py", line 564, in from_pretrained return model_class.from_pretrained( File "/usr/local/lib/python3.9/site-packages/transformers/modeling_utils.py", line 3737, in from_pretrained dtype_orig = cls._set_default_torch_dtype(torch_dtype) File "/usr/local/lib/python3.9/site-packages/transformers/modeling_utils.py", line 1571, in _set_default_torch_dtype torch.set_default_dtype(dtype) File "/usr/local/lib/python3.9/site-packages/torch/init.py", line 796, in set_default_dtype _C._set_default_dtype(d) TypeError: couldn't find storage object Float8_e4m3fnStorage

I set torch_dtype=torch.float8_e4m3fn when loading the llama3.1 405b model using transformers and this error occurs. cuda version is 11.8 and torch version is 2.4.0. transformers version is 4.43.1 Any idea how to fix this error?