dmx.compressor.utils.benchmark.prepare_model

dmx.compressor.utils.benchmark.prepare_model(model: Module, evaluation_mode: EVALUATION_MODE, model_runner: Callable[[Module | DmxModel], None])

Prepares a DMXModel if needed in Baseline, Basic, or Basic_NoVSIMD modes

Parameters:

• model (torch.nn.Module) – torch model to prepare

• evaluation_mode (EVALUATION_MODE) – The mode for the DMXModel to create out of the torch model

• model_runner (Callable[[Union[torch.nn.Module, DmxModel]], None]) – A function for running a sample input through the model