torchgpipe
torchgpipe copied to clipboard
CUDA issues
Hi, I have some questions regarding CUDA usage when I call the GPipe class. When GPipe tries to split_module partition.to(device), the error torch._C._cuda_init() RuntimeError: No CUDA GPUs are available happens. However, if I don't use GPipe under the same environment, the GPUs can be found by torch. I wonder if there is anyone ever faced with the same problem? Thank you!