RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

username: awais19999
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/a/8ab73bd5b77279c479c124e51bc3457879732068.png" data-download-href="https://discuss.pytorch.org/uploads/default/8ab73bd5b77279c479c124e51bc3457879732068" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/8/a/8ab73bd5b77279c479c124e51bc3457879732068.png" alt="image" data-base62-sha1="jN8sYB9vUFS6EFxnnwmKaMvLVeE" width="690" height="156" data-dominant-color="303030"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1357×307 15.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
am getting this error while train the model, could anybody tell me how can i resolve this issue? following is the code</p>
<p>class LinearNoiseScheduler:<br>
r"“”<br>
Class for the linear noise scheduler that is used in DDPM.<br>
“”"<br>
def <strong>init</strong>(self, num_timesteps, beta_start, beta_end):<br>
self.num_timesteps = num_timesteps<br>
self.beta_start = beta_start<br>
self.beta_end = beta_end</p>
<pre><code>    self.betas = torch.linspace(beta_start, beta_end, num_timesteps)
    self.alphas = 1. - self.betas
    self.alpha_cum_prod = torch.cumprod(self.alphas, dim=0)
    self.sqrt_alpha_cum_prod = torch.sqrt(self.alpha_cum_prod)
    self.sqrt_one_minus_alpha_cum_prod = torch.sqrt(1 - self.alpha_cum_prod)
    
def add_noise(self, original, noise, t):
    r"""
    Forward method for diffusion
    :param original: Image on which noise is to be applied
    :param noise: Random Noise Tensor (from normal dist)
    :param t: timestep of the forward process of shape -&gt; (B,)
    :return:
    """
    original_shape = original.shape
    batch_size = original_shape[0]
    
    sqrt_alpha_cum_prod = self.sqrt_alpha_cum_prod[t].reshape(batch_size)
    sqrt_one_minus_alpha_cum_prod = self.sqrt_one_minus_alpha_cum_prod[t].reshape(batch_size)
    
    # Reshape till (B,) becomes (B,1,1,1) if image is (B,C,H,W)
    for _ in range(len(original_shape)-1):
        sqrt_alpha_cum_prod = sqrt_alpha_cum_prod.unsqueeze(-1)
    for _ in range(len(original_shape)-1):
        sqrt_one_minus_alpha_cum_prod = sqrt_one_minus_alpha_cum_prod.unsqueeze(-1)

    # Apply and Return Forward process equation
    return (sqrt_alpha_cum_prod.to(original.device) * original
            + sqrt_one_minus_alpha_cum_prod.to(original.device) * noise)
    
def sample_prev_timestep(self, xt, noise_pred, t):
    r"""
        Use the noise prediction by model to get
        xt-1 using xt and the nosie predicted
    :param xt: current timestep sample
    :param noise_pred: model noise prediction
    :param t: current timestep we are at
    :return:
    """
    x0 = (xt - (self.sqrt_one_minus_alpha_cum_prod[t] * noise_pred)) / torch.sqrt(self.alpha_cum_prod[t])
    x0 = torch.clamp(x0, -1., 1.)

    mean = xt - ((self.betas[t])*noise_pred)/(self.sqrt_one_minus_alpha_cum_prod[t])
    mean = mean / torch.sqrt(self.alphas[t])
    
    if t == 0:
        return mean, mean
    else:
        variance = (1-self.alpha_cum_prod[t-1]) / (1.0 - self.alpha_cum_prod[t])
        variance = variance * self.betas[t]
        sigma = variance ** 0.5
        z = torch.randn(xt.shape).to(xt.device)
        
        # OR
        # variance = self.betas[t]
        # sigma = variance ** 0.5
        # z = torch.randn(xt.shape).to(xt.device)
        return mean + sigma*z, x0
</code></pre>

username: ptrblck
cooked: <p>Based on the error message it seems <code>t</code> might be on the wrong device in:</p>
<pre data-code-wrap="python"><code class="lang-python">sqrt_alpha_cum_prod = self.sqrt_alpha_cum_prod[t]
</code></pre>
<p>Could you check its <code>.device</code> attribute and move it to the same one <code>self.sqrt_alpha_cum_prod</code> uses?</p>