code: # Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #    http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing,
  # software distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import os
  from typing import Union
  
  import pytest
  import torch
  
  from sparseml.pytorch.models import ModelRegistry, mobilenet
  from tests.sparseml.pytorch.models.utils import compare_model
  
  
  @pytest.mark.skipif(
      os.getenv("NM_ML_SKIP_PYTORCH_TESTS", False),
      reason="Skipping pytorch tests",
  )
  @pytest.mark.skipif(
      os.getenv("NM_ML_SKIP_MODEL_TESTS", False),
      reason="Skipping model tests",
  )
  @pytest.mark.parametrize(
      "key,pretrained,test_input",
      [
          ("mobilenet", False, True),
          ("mobilenet", True, False),
          ("mobilenet", "base", False),
          ("mobilenet", "pruned-conservative", False),
          ("mobilenet", "pruned-moderate", False),
      ],
  )
  def test_mobilenets(key: str, pretrained: Union[bool, str], test_input: bool):
      model = ModelRegistry.create(key, pretrained)
      diff_model = mobilenet()
  
      if pretrained:
          compare_model(model, diff_model, same=False)
          match_model = ModelRegistry.create(key, pretrained)
          compare_model(model, match_model, same=True)
  
      if test_input:
          input_shape = ModelRegistry.input_shape(key)
          batch = torch.randn(1, *input_shape)
          out = model(batch)
          assert isinstance(out, tuple)
          for tens in out:
              assert tens.shape[0] == 1
              assert tens.shape[1] == 1000
  
apis: torch.randn