How to load the model i trained and denoise from a given step?

i'm using the training code offered by the author:

from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer

model = Unet( dim = 64, dim_mults = (1, 2, 4, 8), flash_attn = True )

diffusion = GaussianDiffusion( model, image_size = 128, timesteps = 1000, # number of steps sampling_timesteps = 250 # number of sampling timesteps (using ddim for faster inference [see citation for ddim paper]) )

trainer = Trainer( diffusion, 'ImageDataset', train_batch_size = 32, train_lr = 8e-5, train_num_steps = 700000, # total training steps gradient_accumulate_every = 2, # gradient accumulation steps ema_decay = 0.995, # exponential moving average decay amp = True, # turn on mixed precision calculate_fid = True # whether to calculate fid during training )

trainer.train()

Now I'm trying to verify the model I trained in another 'denoising.py' file, let's say the model's save path is "./results/model-best.pt" Should I type the code above again except trainer.train() and add trainer.load("best")? And do that means after this I can just type some code such as

for t in tqdm(reversed(range(0, t_max+1)), desc = 'sampling loop time step', total = t_max+1):

img, x_start = diffusion.p_sample(img, t)

to start denoising from the step t_max and image img i want? Excuse me for my bad expression.