Add missing caffe2_hip extension in setup.py

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/15189

Reviewed By: orionr

Differential Revision: D13457644

Pulled By: bddppq

fbshipit-source-id: c2363e9b8fd21709b62777e5b2199f01ec1c65f8

diff --git a/setup.py b/setup.py
index 6f8323d..79ea3b7 100644 (file)
--- a/setup.py
+++ b/setup.py
@@ -908,6 +908,12 @@ if USE_CUDA:
             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,