aimet_torch.cross_layer_equalization¶

Top-level API

aimet_torch.cross_layer_equalization.equalize_model(model, input_shapes=None, dummy_input=None)[source]¶

High-level API to perform Cross-Layer Equalization (CLE) on the given model. The model is equalized in place.

Parameters:

model (Module) – Model to equalize
input_shapes (Union[Tuple, List[Tuple], None]) – Shape of the input (can be a tuple or a list of tuples if multiple inputs)
dummy_input (Union[Tensor, Tuple, None]) – A dummy input to the model. Can be a Tensor or a Tuple of Tensors. dummy_input will be placed on CPU if not already.