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

[hoalarious]

Getting an error:

Traceback (most recent call last):
  File "/kaggle/working/diffused-heads-colab/sample.py", line 24, in main
    samples = diffusion.sample(id_frame, audio_emb.unsqueeze(0), **args.unet)
  File "/kaggle/working/diffused-heads-colab/diffusion.py", line 42, in sample
    sample_frame = self.sample_loop(xT[:, i].to(x_cond.device), x_cond, motion_frames, audio_emb[:, audio_ids])
  File "/kaggle/working/diffused-heads-colab/diffusion.py", line 54, in sample_loop
    mean, logvar = self.get_p_params(xt, timesteps_ids, nn_out)
  File "/kaggle/working/diffused-heads-colab/diffusion.py", line 67, in get_p_params
    p_logvar = nu * self.expand(torch.log(self.beta[timesteps])) + (1 - nu) * self.expand(self.log_beta_tilde_clipped[timesteps])
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.

Trying to get this working on colab/kaggle.

Need to setup my local environment to debug this better. Wondering if anyone has a fix for this in the mean time?

[Bebaam]

Not sure how it impacts performance, but at least for testing you can just put self.beta[timesteps] and self.log_beta_tilde_clipped to your GPU like this in the line before:
self.beta = self.beta.to(xt.device) and self.log_beta_tilde_clipped = self.log_beta_tilde_clipped.to(xt.device)

The same you have to do for the subsequent errors, then you should be able to successfully sample.

[hoalarious]

Thanks I was hoping for a more elegant solution but probably not a good idea to get tied down on it at this stage. This talking head implementation is rather computational intensive. Taking 16 minutes for a 2 second clip. Going to test it on longer outputs to see if it gets better. I don't think that's from the fixes you suggested.

Got a working colab demo if anyone wants to try it:
https://github.com/hoalarious/diffused-heads-colab/blob/main/diffused_heads_colab.ipynb

[akarinmoe]

Not sure how it impacts performance, but at least for testing you can just put self.beta[timesteps] and self.log_beta_tilde_clipped to your GPU like this in the line before: self.beta = self.beta.to(xt.device) and self.log_beta_tilde_clipped = self.log_beta_tilde_clipped.to(xt.device)

The same you have to do for the subsequent errors, then you should be able to successfully sample.

Thanks a lot
I try what you do and modify the function get_q_params as well

def get_q_params(self, xt, timesteps, eps_pred=None, x0=None):
        if x0 is None:
        # predict x0 from xt and eps_pred
            self.coef1_x0 = self.coef1_x0.to(xt.device)
            coef1_x0 = self.expand(self.coef1_x0[timesteps])
            self.coef2_x0 = self.coef2_x0.to(xt.device)
            coef2_x0 = self.expand(self.coef2_x0[timesteps])
            x0 = coef1_x0 * xt - coef2_x0 * eps_pred
            x0 = x0.clamp(-1, 1)

        # q(x_{t-1} | x_t, x_0)
        self.coef1_q = self.coef1_q.to(xt.device)
        coef1_q = self.expand(self.coef1_q[timesteps])
        self.coef2_q = self.coef2_q.to(xt.device)
        coef2_q = self.expand(self.coef2_q[timesteps])
        q_mean = coef1_q * x0 + coef2_q * xt

        q_logvar = self.expand(self.log_beta_tilde_clipped[timesteps])

        return q_mean, q_logvar

I move every param to the device GPU in this function and finally solve the problem