Getting nan as generated image by using cosine schedule for DDPMScheduler #468

model = DiffusionModelUNet(
 spatial_dims=2,
 in_channels=1,
 out_channels=1,
 num_channels=(64,128, 256,512,1024),
 attention_levels=(False,False, True, True,True),
 num_res_blocks=3,
 num_head_channels=256,
)
model.to(device)
scheduler = DDPMScheduler(num_train_timesteps=total_timesteps,schedule="cosine")
optimizer = torch.optim.Adam(params=model.parameters(), lr=2.5e-5)
inferer = DiffusionInferer(scheduler)

if use_pretrained_model==False:
 epoch_loss_list = []
 val_epoch_loss_list = []
 memory_list = []
 scaler = GradScaler()
 total_start = time.time()
 for epoch in range(n_epochs):
 #print(f"memory usage before epoch {epoch} is {torch.cuda.mem_get_info(device)}")
 mem = torch.cuda.mem_get_info(device)
 frac = (1-(mem[0]/1000000000)/(mem[1]/1000000000))*100
 print(f"fraction gpu memory usage before epoch {epoch} is {frac:.2f}%")
 memory_list.append(frac)
 model.train()
 epoch_loss = 0
 progress_bar = tqdm(enumerate(train_loader), total=len(train_loader), ncols=70)
 progress_bar.set_description(f"Epoch {epoch}")
 for step, batch in progress_bar:
 batch_ct, batch_organmap = batch
 images = batch_ct.to(device) 
 optimizer.zero_grad(set_to_none=True)
 with autocast(enabled=True):
 # Generate random noise
 noise = torch.randn_like(images).to(device)
 # Create timesteps
 timesteps = torch.randint(
 0, inferer.scheduler.num_train_timesteps, (images.shape[0],), device=images.device
 ).long()
 # Get model prediction
 noise_pred = inferer(inputs=images, diffusion_model=model, noise=noise, timesteps=timesteps)
 loss = F.mse_loss(noise_pred.float(), noise.float())
 scaler.scale(loss).backward()
 scaler.step(optimizer)
 scaler.update()
 epoch_loss += loss.item()
 progress_bar.set_postfix({"loss": epoch_loss / (step + 1)})
 epoch_loss_list.append(epoch_loss / (step + 1))
 if (epoch + 1) % val_interval == 0:
 model.eval()
 val_epoch_loss = 0
 for step, batch in enumerate(test_loader):
 batch_ct, batch_organmap = batch
 images = batch_ct.to(device) 
 with torch.no_grad():
 with autocast(enabled=True):
 noise = torch.randn_like(images).to(device)
 timesteps = torch.randint(
 0, inferer.scheduler.num_train_timesteps, (images.shape[0],), device=images.device
 ).long()
 noise_pred = inferer(inputs=images, diffusion_model=model, noise=noise, timesteps=timesteps)
 val_loss = F.mse_loss(noise_pred.float(), noise.float())
 val_epoch_loss += val_loss.item()
 progress_bar.set_postfix({"val_loss": val_epoch_loss / (step + 1)})
 val_epoch_loss_list.append(val_epoch_loss / (step + 1))
 # Sampling image during training
 noise = torch.randn((1, 1, 256, 256))
 noise = noise.to(device)
 scheduler.set_timesteps(num_inference_steps=1000)
 with autocast(enabled=True):
 image = inferer.sample(input_noise=noise, diffusion_model=model, scheduler=scheduler)
 print(image.shape)
 print(image)
 plt.figure(figsize=(2, 2))
 plt.imshow(image[0, 0].cpu(), cmap="gray")
 plt.tight_layout()
 plt.axis("off")
 plt.show()
 total_time = time.time() - total_start
 print(f"train completed, total time: {total_time}.")