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?


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:

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: