diff --git a/op_builder/builder.py b/op_builder/builder.py
index 52f1a9cf9..2e3728397 100644
--- a/op_builder/builder.py
+++ b/op_builder/builder.py
@@ -3,14 +3,18 @@ import re
 from pathlib import Path
 from typing import List
 
-import torch
-
 
 def get_cuda_cc_flag() -> List:
     """get_cuda_cc_flag
 
     cc flag for your GPU arch
     """
+
+    # only import torch when needed
+    # this is to avoid importing torch when building on a machine without torch pre-installed
+    # one case is to build wheel for pypi release
+    import torch
+    
     cc_flag = []
     for arch in torch.cuda.get_arch_list():
         res = re.search(r'sm_(\d+)', arch)
diff --git a/setup.py b/setup.py
index 453f6421d..dda1061d7 100644
--- a/setup.py
+++ b/setup.py
@@ -15,8 +15,9 @@ try:
     if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 10):
         raise RuntimeError("Colossal-AI requires Pytorch 1.10 or newer.\n"
                            "The latest stable release can be obtained from https://pytorch.org/")
+    TORCH_AVAILABLE = True
 except ImportError:
-    raise ModuleNotFoundError('torch is not found. You need to install PyTorch before installing Colossal-AI.')
+    TORCH_AVAILABLE = False
 
 
 # ninja build does not work unless include_dirs are abs path
@@ -24,7 +25,7 @@ this_dir = os.path.dirname(os.path.abspath(__file__))
 build_cuda_ext = True
 ext_modules = []
 
-if int(os.environ.get('NO_CUDA_EXT', '0')) == 1:
+if int(os.environ.get('NO_CUDA_EXT', '0')) == 1 or not TORCH_AVAILABLE:
     build_cuda_ext = False