search

pytorch / extension-cpp

C++ extensions in PyTorch
1.02k stars 214 forks source link

How to automatically set the backends? #9

Closed jxgu1016 closed 6 years ago

jxgu1016 commented 6 years ago

Is there any way to automatically set the backends (cpu or gpu)? To merge the two Function into one and it will choose the proper backend according to what device we are using.

jxgu1016 commented 6 years ago

Also, can we merge the two setup.py together and make them share the same package name? Sorry to close this issue mistakenly ...

goldsborough commented 6 years ago

You can merge the two Functions and write:

if tensor.is_cuda():
  # call CUDA code
else:
  # call C++ code

inside your PyTorch Function, or you can move it into C++ and write:

if (tensor.is_cuda()) {
  // call CUDA code
} else {
  // call C++ code
}

There's nothing to stop you from merging the setup.pys, no. I won't do it for this example because things are supposed to be isolated.