diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml
index 4f60e781..e3be0d17 100644
--- a/.github/workflows/unit-test.yml
+++ b/.github/workflows/unit-test.yml
@@ -53,7 +53,7 @@ jobs:
             fi
          fi
       - name: Install TinyNeuralNetwork
-        run: python setup.py install
+        run: pip install .
       - name: Install test dependencies
         env:
           PYTORCH_VER: ${{ matrix.vers.pt_ver }}
@@ -105,7 +105,7 @@ jobs:
       - name: Install PyTorch
         run: conda install pytorch torchvision cpuonly -c pytorch
       - name: Install TinyNeuralNetwork
-        run: python setup.py install
+        run: pip install .
       - name: Check op matrix
         run: |
           python scripts/gen_op_docs.py
diff --git a/tinynn/graph/quantization/quantizer.py b/tinynn/graph/quantization/quantizer.py
index 38219479..a2a64207 100644
--- a/tinynn/graph/quantization/quantizer.py
+++ b/tinynn/graph/quantization/quantizer.py
@@ -60,8 +60,10 @@
 
 try:
     import ruamel_yaml as yaml
+    from ruamel_yaml import CommentedMap
 except ModuleNotFoundError:
     import ruamel.yaml as yaml
+    from ruamel.yaml import CommentedMap
 
 
 # Fusable OPs for Quantize Aware Training
@@ -363,7 +365,7 @@ def __init__(self, model, dummy_input, work_dir: typing.Optional[str] = None, co
 
         self.train_mode_dict = {}
 
-        self.layerwise_config = yaml.comments.CommentedMap()
+        self.layerwise_config = CommentedMap()
         self.effective_layers = []
         self.layerwise_default = True
         if config is not None and 'layerwise_config' in config: