Refactor FastGRNN CUDA Setup

Author: ShikharJ

pytorch/README.md

@@ -68,6 +68,7 @@ Install appropriate CUDA and cuDNN [Tested with >= CUDA 8.1 and cuDNN >= 6.1]
 ```
 pip install -r requirements-gpu.txt
 pip install -e .
+pip install -e edgeml_pytorch/cuda/
 ```
 
 **Note**: For using the optimized FastGRNNCUDA implementation, it is recommended to use CUDA v10.1, gcc 7.5 and cuDNN v7.6 and torch==1.4.0. Also, there are some known issues when compiling custom CUDA kernels on Windows [pytorch/#11004](https://github.com/pytorch/pytorch/issues/11004).

pytorch/edgeml_pytorch/cuda/setup.py

@@ -0,0 +1,18 @@
+import setuptools #enables develop
+import os
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+from edgeml_pytorch.utils import findCUDA
+
+if findCUDA() is not None:
+    setuptools.setup(
+        name='fastgrnn_cuda',
+        ext_modules=[
+            CUDAExtension('fastgrnn_cuda', [
+                'edgeml_pytorch/cuda/fastgrnn_cuda.cpp',
+                'edgeml_pytorch/cuda/fastgrnn_cuda_kernel.cu',
+            ]),
+        ],
+        cmdclass={
+            'build_ext': BuildExtension
+        }
+    )

pytorch/edgeml_pytorch/graph/rnn.py

@@ -9,8 +9,11 @@
 
 import edgeml_pytorch.utils as utils
 
-if utils.findCUDA() is not None:
-    import fastgrnn_cuda
+try:
+    if utils.findCUDA() is not None:
+        import fastgrnn_cuda
+except:
+    pass
 
 
 # All the matrix vector computations of the form Wx are done 

pytorch/setup.py

@@ -1,21 +1,5 @@
 import setuptools #enables develop
 import os
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-from edgeml_pytorch.utils import findCUDA
-
-if findCUDA() is not None:
-    setuptools.setup(
-        name='fastgrnn_cuda',
-        ext_modules=[
-            CUDAExtension('fastgrnn_cuda', [
-                'edgeml_pytorch/cuda/fastgrnn_cuda.cpp',
-                'edgeml_pytorch/cuda/fastgrnn_cuda_kernel.cu',
-            ]),
-        ],
-        cmdclass={
-            'build_ext': BuildExtension
-        }
-    )
 
 setuptools.setup(
     name='edgeml',