From bdfff2f8c28faa71e8323ee5cb8a4957679e581c Mon Sep 17 00:00:00 2001
From: Junjie Bai <bai@in.tum.de>
Date: Thu, 13 Dec 2018 15:57:20 -0800
Subject: [PATCH] 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
---
 setup.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/setup.py b/setup.py
index 6f8323d976019..79ea3b7950c5a 100644
--- 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,