Fix OpenMP flags

ghstack-source-id: d1f59bb
Pull Request resolved: #30

Author: pbelevich

setup.py

@@ -3,15 +3,19 @@
 
 build_cuda = cpp_extension.CUDA_HOME != None
 
+CXX_FLAGS = ['-fopenmp']
+NVCC_FLAGS = ['--expt-extended-lambda', '-Xcompiler', '-fopenmp']
+
 if build_cuda:
     csprng_ext = cpp_extension.CUDAExtension(
         'torch_csprng', ['csprng.cu'],
         extra_compile_args={'cxx': [],
-                            'nvcc': ['-O2', '--expt-extended-lambda']}
+                            'nvcc': NVCC_FLAGS}
     )
 else:
     csprng_ext = cpp_extension.CppExtension(
-        'torch_csprng', ['csprng.cpp']
+        'torch_csprng', ['csprng.cpp'],
+        extra_compile_args={'cxx': CXX_FLAGS}
     )
 
 setup(name='pytorch_csprng',