NaN in DPMSolverMultistepInverseScheduler #10748

dpmpp_2m_sde_karras_scheduler_inv = DPMSolverMultistepInverseScheduler(
    num_train_timesteps=1000,
    beta_start=0.00085,
    beta_end=0.012,
    algorithm_type="sde-dpmsolver++",
    use_karras_sigmas=True,
    steps_offset=1
)

self.scheduler.set_timesteps(self.inv_config.steps)
timesteps = self.scheduler.timesteps
with torch.autocast(device_type=self.device, dtype=self.dtype):
    for i, t in enumerate(tqdm(timesteps)):
        noises = []
        x_index = torch.arange(len(x))
        batches = x_index.split(self.batch_size, dim = 0)
        for batch in batches:
            noise = self.pred_noise(
                x[batch], conds, timesteps[i], concat_conds=x[batch], batch_idx=batch)
            noises += [noise]
        noises = torch.cat(noises)
        
        x = self.scheduler.step(noises, t, x, generator=self.rng, return_dict=False)[0]