When a tensor is generated from some_func(A.shape) (where A is a tensor), the generated tensor locates in cpu, not A's device

how to solve it ? I have tried tensor.to(A.device) and tensor.to(accelerator.device), but it seems not to work.

Hi @weizhenhuan , can you provide a reproducer ?