get_custom_module_class_keys#
- class torch.ao.quantization.fx.utils.get_custom_module_class_keys(custom_module_mapping)[source]#
Get all the unique custom module keys in the custom config dict.
Example input:
{ QuantType.STATIC: {CustomModule1: ObservedCustomModule}, QuantType.DYNAMIC: {CustomModule2: DynamicObservedCustomModule}, QuantType.WEIGHT_ONLY: {CustomModule3: WeightOnlyObservedCustomModule}, }
Example output:
# extract the keys across all inner STATIC, DYNAMIC, and WEIGHT_ONLY dicts [CustomModule1, CustomModule2, CustomModule3]