ValueError: Trying to set a tensor of shape torch.Size([128256, 1024]) in "weight"
File "/home/rkuo/GenAI/Text-to-Text/llm_server.py", line 58, in
LLM = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="cuda")
File "/home/rkuo/.local/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 563, in from_pretrained
return model_class.from_pretrained(
File "/home/rkuo/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 3677, in from_pretrained
) = cls._load_pretrained_model(
File "/home/rkuo/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4104, in _load_pretrained_model
new_error_msgs, offload_index, state_dict_index = _load_state_dict_into_meta_model(
File "/home/rkuo/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 886, in _load_state_dict_into_meta_model
set_module_tensor_to_device(model, param_name, param_device, **set_module_kwargs)
File "/home/rkuo/.local/lib/python3.10/site-packages/accelerate/utils/modeling.py", line 358, in set_module_tensor_to_device
raise ValueError(
ValueError: Trying to set a tensor of shape torch.Size([128256, 1024]) in "weight" (which has shape torch.Size([128256, 4096])), this look incorrect.