Hi,
I trained a model defined using custom class on 8-GPU setup using Accelerate library. The model trains on multi-GPU setup and is saved successfully.
Saving model:
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.eval()
accelerator.save(unwrapped_model.state_dict(), save_model)

Now I want to fine-tune previously saved model on different dataset and this is how I load the model:
model.load_state_dict(torch.load(model_path_level_1))

And then pass it on to accelerate.prepare:
    accelerator = Accelerator(kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=True)]) 
model, optimizer, train_dataloader, validation_dataloader = accelerator.prepare(model, optimizer, train_dataloader, validation_dataloader)

This script launches 8 processes on GPU:0 and the code fails with Out of memory issues after sometime. After launching 8 processes, one process each is launched on other 7 GPUs before script crashes.
Is this a problem with loading custom class models?
Thank you.