pimm.export.clean_state_dict

clean_state_dict(state_dict: Dict[str, Any]) → Dict[str, Any]

Remove DDP and torch.compile prefixes from state dict keys.

This ensures compatibility between single-GPU and multi-GPU checkpoints.

Parameters:

state_dict – State dictionary with potentially prefixed keys.

Returns:

Cleaned state dictionary.