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 
                                             else torch.FloatTensor)

Now, use this as a switch.
use_gpu()
''' ...
your code here
...'''

Thank you team for making life so simple