How to load the saved model?

[bilalghanem]

Hi

I need your help with loading the model. I see how you're doing that in the "converting..." file. But this is only for LORA models.
What about full_shard models? (--sharding_strategy full_shard --train_type full).

I tried to load it this way but it didn't work:

model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")
model.load_state_dict(torch.load('model_state_dict.safetensors'))