From 22fbaf44396d89064688ff570db77393a0b7a8a6 Mon Sep 17 00:00:00 2001
From: Ean Garvey <87458719+monorimet@users.noreply.github.com>
Date: Thu, 4 Apr 2024 09:55:38 -0500
Subject: [PATCH] Workaround data dependent control flow issues in PNDM

---
 src/diffusers/schedulers/scheduling_pndm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/diffusers/schedulers/scheduling_pndm.py b/src/diffusers/schedulers/scheduling_pndm.py
index 4bfd9d6c3a7f8..4970581f93645 100644
--- a/src/diffusers/schedulers/scheduling_pndm.py
+++ b/src/diffusers/schedulers/scheduling_pndm.py
@@ -418,7 +418,7 @@ def _get_prev_sample(self, sample, timestep, prev_timestep, model_output):
         # model_output -> e_θ(x_t, t)
         # prev_sample -> x_(t−δ)
         alpha_prod_t = self.alphas_cumprod[timestep]
-        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
+        torch.where(prev_timestep >= 0, torch.index_select(self.alphas_cumprod, 0, prev_timestep), self.final_alpha_cumprod)
         beta_prod_t = 1 - alpha_prod_t
         beta_prod_t_prev = 1 - alpha_prod_t_prev