diff --git a/birefnet.py b/birefnet.py
index 1ed28de..d68d7d3 100644
--- a/birefnet.py
+++ b/birefnet.py
@@ -2,9 +2,9 @@
 
 import os
 import math
+from transformers import PretrainedConfig
 
-
-class Config():
+class Config(PretrainedConfig):
     def __init__(self) -> None:
         # PATH settings
         self.sys_home_dir = os.path.expanduser('~')    # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx