Update README.md

README.md

@@ -30,7 +30,7 @@ Next use the following straightforward steps to generate images from the final m
 ```
 from micro_diffusion.models.model import create_latent_diffusion
 model = create_latent_diffusion(latent_res=64, in_channels=4, pos_interp_scale=2.0).to('cuda')
-model.load_state_dict(torch.load(ckpt_path_on_local_disk))
+model.dit.load_state_dict(torch.load(ckpt_path_on_local_disk))
 gen_images = model.generate(prompt=['An elegant squirrel pirate on a ship']*4, num_inference_steps=30, 
                                     guidance_scale=5.0, seed=2024)
 ```