How to specify GPU usage?

10

How does one use GPUs if one has a custom NN class (that inherits from torch.nn.Module)?

For example, I know that using the easy example from (http://pytorch.org/tutorials/beginner/pytorch_with_examples.html) one can just change the type of the tensors being created:

dtype = torch.FloatTensor
# dtype = torch.cuda.FloatTensor # Uncomment this to run on GPU

however, when using things like torch.nn.Linear and also Variable, how does one make sure to use GPUs?

Also, do I really have to track how GPUs are assigned, I am fine with torch just doing its stuff automagically.

In particular I would love to see how:

http://pytorch.org/tutorials/beginner/examples_nn/two_layer_net_module.html#sphx-glr-beginner-examples-nn-two-layer-net-module-py

is made into a GPU version of it.

Related SO question: https://stackoverflow.com/questions/45553613/how-does-one-make-sure-that-everything-is-running-on-gpu-automatically-in-pytorc