diff --git a/main.py b/main.py
index 2570022..3a7e348 100644
--- a/main.py
+++ b/main.py
@@ -4,7 +4,7 @@
 from tensorflow.python.saved_model import signature_constants
 from tensorflow.python.saved_model import tag_constants
 import os
-from config import mask_config
+from config import mask_config, CocoConfig
 from model import MaskRCNN
 
 sess = tf.Session()
@@ -13,8 +13,7 @@
 
 def get_config():
     if is_coco:
-        import coco
-        class InferenceConfig(coco.CocoConfig):
+        class InferenceConfig(CocoConfig):
             GPU_COUNT = 1
             IMAGES_PER_GPU = 1
 
@@ -115,4 +114,4 @@ def make_serving_ready(model_path, save_serve_path, version_number):
                      PATH_TO_SAVE_TENSORFLOW_SERVING_MODEL,
                      VERSION_NUMBER)
 
-print("COMPLETED")
\ No newline at end of file
+print("COMPLETED")