1# Copyright (c) Meta Platforms, Inc. and affiliates. 2# All rights reserved. 3# 4# This source code is licensed under the BSD-style license found in the 5# LICENSE file in the root directory of this source tree. 6 7import logging 8 9import torch 10 11from torchvision import models 12 13from ..model_base import EagerModelBase 14 15 16class MV3Model(EagerModelBase): 17 def __init__(self): 18 pass 19 20 def get_eager_model(self) -> torch.nn.Module: 21 logging.info("Loading mobilenet_v3 model") 22 mv3_small = models.mobilenet_v3_small( 23 weights=models.MobileNet_V3_Small_Weights.IMAGENET1K_V1 24 ) 25 logging.info("Loaded mobilenet_v3 model") 26 return mv3_small 27 28 def get_example_inputs(self): 29 tensor_size = (1, 3, 224, 224) 30 return (torch.randn(tensor_size),) 31