In older versions of PyTorch, in order to move everything to the GPU, one had to do the following.
# Define a lambda at the top
cuda = lambda x: x.cuda() if torch.cuda.is_available() else x
x = Variable(cuda(torch.randn(10))) # When creating variables
model = cuda(Model()) # When creating modules
With the release of PyTorch 0.4, this has been slightly simplified as:
# Define the default tensor type at the top
torch.set_default_tensor_type(torch.cuda.FloatTensor if torch.cuda.is_available()
# When creating tensors
x = torch.randn(10) # No need for additional bells and whistles
However, while creating modules, there is no such simplification:
device = torch.device('cuda' if torch.cuda.is_avaliable() else 'cpu')
model = Model().to(device) # One has to still cast the module to the appropriate device
Is there any workaround possible?
Edit: This method doesn’t seem to work for reasons mentioned in a later post.
Just realized that you don’t actually need to cast modules to the device.
Setting the default tensor type applies to all parameters and tensors automatically within the module.
So, my current method for running a script on GPU is the following:
# Define this lambda
use_gpu = lambda x=True: torch.set_default_tensor_type(torch.cuda.FloatTensor
if torch.cuda.is_available() and x
Now, use this as a switch.
your code here
Thank you team for making life so simple
Sorry to have revived this thread again.
Just tried it out and it seems that there are some fundamental problems with setting the default tensor type.
Dataloaders give normal (non-cuda) tensors by default. They have to be manually cast using the
Many methods are simply not implemented for
torch.cuda.*Tensor. Thus, setting the global tensor type to cuda fails.
Conversions to numpy using the
numpy() method aren’t’ available for cuda tensors. One has to go
Although this chain is agnostic, it defeats the purpose.
Can someone from the team pitch in on this?
Is there any hope for a global device setting that just works?