Add missing caffe2_hip extension in setup.py

Summary: Pull Request resolved: pytorch#15189 Reviewed By: orionr Differential Revision: D13457644 Pulled By: bddppq fbshipit-source-id: c2363e9b8fd21709b62777e5b2199f01ec1c65f8
harouwu · Dec 13, 2018 · bdfff2f · bdfff2f
1 parent de07845
commit bdfff2f
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/setup.py b/setup.py
@@ -908,6 +908,12 @@ def make_relative_rpath(path):
             name=str('caffe2.python.caffe2_pybind11_state_gpu'),
             sources=[]),
     )
+if USE_ROCM:
+    extensions.append(
+        Extension(
+            name=str('caffe2.python.caffe2_pybind11_state_hip'),
+            sources=[]),
+    )
 
 cmdclass = {
     'create_version_file': create_version_file,