Resolving torch cuda availablity for cpu-only installation

setup.py

@@ -16,6 +16,7 @@ import os
 from setuptools import setup, Extension, find_packages
 import subprocess
 
+import torch
 from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension, CUDA_HOME
 
 from scripts.utils import get_nvidia_cc
@@ -37,7 +38,7 @@ extra_cuda_flags = [
 ]
 
 def get_cuda_bare_metal_version(cuda_dir):
-    if cuda_dir==None:
+    if cuda_dir==None or torch.version.cuda==None:
         print("CUDA is not found, cpu version is installed")
         return None, -1, 0
     else: