How to move data or model from CPU to GPU?

Jittor / jittor

Jittor is a high-performance deep learning framework based on JIT compiling and meta-operators.

https://cg.cs.tsinghua.edu.cn/jittor/

Apache License 2.0

3.1k stars 313 forks source link

How to move data or model from CPU to GPU? #394

Open HeroesHanzo opened 2 years ago

HeroesHanzo commented 2 years ago

Dear the developers of Jittor. I'm novice in Jittor and I used to adopt Pytorch in my deep learning project. When I wanna move my data or model from CPU to GPU, I always apply .to('cuda') operation. Is there any similar function in Jittor? If not, is there any explicit operation to move the data or models from CPU to GPU? Or I just need to specify the GPU (CUDA_VISIBLE_DEVICE=XXX) before running my script without explicit movement?

I'm looking forward to your reply. Thanks for your help.

Exusial commented 2 years ago

Hello, using jt.flags.use_cuda=1 will put all arrays after this declaration to GPU.