ds_report bug fix on cpu and guard torch import in setup.py (#524)

* on cpu box error gracefully if cuda home doesn't exist

* gaurd against torch import issue

* fix sytax error

* fix import

deepspeed/env_report.py

@@ -62,6 +62,8 @@ def ninja_installed():
 def nvcc_version():
     import torch.utils.cpp_extension
     cuda_home = torch.utils.cpp_extension.CUDA_HOME
+    if cuda_home is None:
+        return f"{RED} [FAIL] cannot find CUDA_HOME via torch.utils.cpp_extension.CUDA_HOME={torch.utils.cpp_extension.CUDA_HOME} {END}"
     try:
         output = subprocess.check_output([cuda_home + "/bin/nvcc",
                                           "-V"],

setup.py

@@ -9,12 +9,17 @@ The wheel will be located at: dist/*.whl
 """
 
 import os
-import torch
 import shutil
 import subprocess
 import warnings
 from setuptools import setup, find_packages
-from torch.utils.cpp_extension import CUDAExtension, BuildExtension, CppExtension
+
+try:
+    import torch
+    from torch.utils.cpp_extension import BuildExtension
+except ImportError:
+    raise ImportError('Unable to import torch, please visit https://pytorch.org/ '
+                      'to see how to properly install torch on your system.')
 
 import op_builder