USING PI0 WARNING: skipping cudagraphs due to skipping cudagraphs due to cpu device (cat).

[Oooer8]

Ticket Type

🐛 Bug Report (Something isn't working)

Environment & System Info

lerobot==0.4.4

Description

when i use lerobot-train CLI to train pi0, i encounter a log as follows:

skipping cudagraphs due to skipping cudagraphs due to cpu device (cat). Found from : 
   File "/home/jiajun/.local/lib/python3.10/site-packages/lerobot/policies/pi0/modeling_pi0.py", line 751, in forward
    time = self.sample_time(actions.shape[0], actions.device)
  File "/home/jiajun/.local/lib/python3.10/site-packages/lerobot/policies/pi0/modeling_pi0.py", line 631, in sample_time
    time_beta = sample_beta(
  File "/home/jiajun/.local/lib/python3.10/site-packages/lerobot/policies/pi0/modeling_pi0.py", line 100, in sample_beta
    return dist.sample((bsize,)).to(device)
  File "/ssd/jiajun/miniconda3/envs/gloveumi/lib/python3.10/site-packages/torch/distributions/distribution.py", line 168, in sample
    return self.rsample(sample_shape)
  File "/ssd/jiajun/miniconda3/envs/gloveumi/lib/python3.10/site-packages/torch/distributions/beta.py", line 78, in rsample
    return self._dirichlet.rsample(sample_shape).select(-1, 0)
  File "/ssd/jiajun/miniconda3/envs/gloveumi/lib/python3.10/site-packages/torch/distributions/dirichlet.py", line 78, in rsample
    concentration = self.concentration.expand(shape)

Context & Reproduction

No response

Relevant logs or stack trace

Checklist

• I have searched existing tickets to ensure this isn't a duplicate.
• I am using the latest version of the main branch.
• I have verified this is not an environment-specific problem.

Additional Info / Workarounds

No response