mmselfsup.utils.misc 源代码
# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn
from mmengine.model import BaseModel, is_model_wrapper
[文档]def get_model(model: nn.Module) -> BaseModel:
"""Get model if the input model is a model wrapper.
Args:
model (nn.Module): A model may be a model wrapper.
Returns:
BaseModel: The model without model wrapper.
"""
if is_model_wrapper(model):
return model.module
else:
return model