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.
2024-01-24T15:13:12Z
Hi@weizhenhuan, can you provide a reproducer ?