Hi ,friends.I want to know that PyTorch how to use cuda


(mshmoon) #1

Recently,I am learning GPU programming, so I know cuda a little. But I have not see that PyTorch how to use the cuda.So ,is there someone help me?


#2

If you just want to use the GPU to train your model etc., you can simply push the data and model to your device using:

model.to('cuda')
data = data.to('cuda')
output = model(data)  # will be performed on the GPU

You don’t need to write CUDA code for it or would you like to implement a specific CUDA kernel for a method?
In that case, have a look at the GitHub issues where you might find a good starter project. :wink: