CollosalAIModelWrapper

class mmengine._strategy.colossalai.CollosalAIModelWrapper(model_wrapper, model)

Parameters

• model_wrapper (None) –

• model (torch.nn.modules.module.Module) –

test_step(data)

Gets the prediction of module during validation process.

Parameters

data (dict or tuple or list) – Data sampled from dataset.

Returns

The predictions of given data.

Return type

list

val_step(data)

Gets the prediction of module during validation process.

Parameters

data (dict or tuple or list) – Data sampled from dataset.

Returns

The predictions of given data.

Return type

list