def METHOD_NAME(self, state_dict, strict=True, model_cfg=None):
    if self.use_sharded_state:
        return super().load_local_state_dict(state_dict, strict=strict)
    else:
        state_dict = dist_utils.broadcast_object(
            state_dict, src_rank=0, group=self.process_group
        )
        return super().METHOD_NAME(state_dict, strict=strict)