evaluation during training causes the stucking

[garychan22]

Hi, thanks for your excellent project !

I want to evaluate the model during training by adding the following code in the train.py script, which leads to program stucking.

if i % conf.evaluate.valid_every == 0:
    model.eval()
    img = diffusion.p_sample_loop(model, [16, 3, 256, 256], model.device)
    save_image(img, "sample.png", normalize=True, range=(-1,1), nrow=4)
    model.train()

Appreciate any help !