I do not think you can specify that you want to use cuda tensors by default.
However you should have a look to the pytorch offical examples.
In the imagenet training/testing script, they use a wrapper over the model called DataParallel.
This wrapper has two advantages:
- it handles the data parallelism over multiple GPUs
- it handles the casting of cpu tensors to cuda tensors
As you can see in L164, you don't have to cast manually your inputs/targets to cuda.
Note that, if you have multiple GPUs and you want to use a single one, launch any python/pytorch scripts with the CUDA_VISIBLE_DEVICES prefix. For instance CUDA_VISIBLE_DEVICES=0 python main.py
.
与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…