Composite Config cannot be saved?

I am trying to create a PretrainedConfig, which uses another config as sub-config. However, the resulting composite config cannot be saved via save_pretrained method and thus causes the Trainer to crash. Anyone can tell me what i did wrong? Sample code for quick replication:

from transformers import PretrainedConfig

class ConfigA(PretrainedConfig):
    def __init__(self, a=10):
        self.a = a

class ConfigB(PretrainedConfig):
    def __init__(self, b=1):
        self.b = b

class CompositeConfig(PretrainedConfig):
    is_composition = True

    def __init__(self, a: ConfigA, b: ConfigB, c=1):
        self.a = a
        self.b = b
        self.c = c

a = ConfigA()
b = ConfigB()
c = CompositeConfig(a, b)



Traceback (most recent call last):
  File "/Users/andre/PycharmProjects/aho-master-ocr/tests/", line 32, in <module>
  File "/Users/andre/opt/anaconda3/envs/easyocr/lib/python3.8/site-packages/transformers/", line 457, in save_pretrained
    self.to_json_file(output_config_file, use_diff=True)
  File "/Users/andre/opt/anaconda3/envs/easyocr/lib/python3.8/site-packages/transformers/", line 900, in to_json_file
  File "/Users/andre/opt/anaconda3/envs/easyocr/lib/python3.8/site-packages/transformers/", line 886, in to_json_string
    return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
  File "/Users/andre/opt/anaconda3/envs/easyocr/lib/python3.8/json/", line 234, in dumps
    return cls(
  File "/Users/andre/opt/anaconda3/envs/easyocr/lib/python3.8/json/", line 201, in encode
    chunks = list(chunks)
  File "/Users/andre/opt/anaconda3/envs/easyocr/lib/python3.8/json/", line 431, in _iterencode
    yield from _iterencode_dict(o, _current_indent_level)
  File "/Users/andre/opt/anaconda3/envs/easyocr/lib/python3.8/json/", line 405, in _iterencode_dict
    yield from chunks
  File "/Users/andre/opt/anaconda3/envs/easyocr/lib/python3.8/json/", line 438, in _iterencode
    o = _default(o)
  File "/Users/andre/opt/anaconda3/envs/easyocr/lib/python3.8/json/", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type ConfigA is not JSON serializable

After looking at the EncoderDecoderConfig it seems to me that is_composition doesn’t really handle anything for us regarding sub-configs and we need to overwrite to_dict and handle the Config Parameters ourselves. Rewriting CompositeConfig in the following way seems to fix the issue:

class CompositeConfig(PretrainedConfig):
    is_composition = True

    def __init__(self, a: Union[Dict, ConfigA], b: Union[Dict, ConfigB], c=1):
        self.a = a if isinstance(a, PretrainedConfig) else ConfigA(**a)
        self.b = b if isinstance(b, PretrainedConfig) else ConfigB(**b)
        self.c = c

    def to_dict(self):
        output = copy.deepcopy(self.__dict__)
        for k, v in output.items():
            if isinstance(v, PretrainedConfig):
                output[k] = v.to_dict()
        output["model_type"] = self.__class__.model_type
        return output