def run(
    argv=None,
    model_class=None,
    model_params=None,
    save_main_session=True,
    device='CPU',
    test_pipeline=None) -> PipelineResult:

  known_args, pipeline_args = parse_known_args(argv)
  pipeline_options = PipelineOptions(pipeline_args)
  pipeline_options.view_as(SetupOptions).save_main_session = save_main_session

  if not model_class:
    # default model class will be mobilenet with pretrained weights.
    model_class = models.mobilenet_v2
    model_params = {'num_classes': 1000}

  def preprocess(image_name: str) -> Tuple[str, torch.Tensor]:
    image_name, image = read_image(
      image_file_name=image_name,
      path_to_dir=known_args.images_dir)
    return (image_name, preprocess_image(image))

  def postprocess(element: Tuple[str, PredictionResult]) -> str:
    filename, prediction_result = element
    prediction = torch.argmax(prediction_result.inference, dim=0)
    return filename + ',' + str(prediction.item())

  # In this example we pass keyed inputs to RunInference transform.
  # Therefore, we use KeyedModelHandler wrapper over PytorchModelHandler.
  model_handler = KeyedModelHandler(
      PytorchModelHandlerTensor(
          state_dict_path=known_args.model_state_dict_path,
          model_class=model_class,
          model_params=model_params,
          device=device,
          min_batch_size=10,
          max_batch_size=100)).with_preprocess_fn(
              preprocess).with_postprocess_fn(postprocess)

  pipeline = test_pipeline
  if not test_pipeline:
    pipeline = beam.Pipeline(options=pipeline_options)

  filename_value_pair = (
      pipeline
      | 'ReadImageNames' >> beam.io.ReadFromText(known_args.input)
      | 'FilterEmptyLines' >> beam.ParDo(filter_empty_lines))
  predictions = (
      filename_value_pair
      | 'PyTorchRunInference' >> RunInference(model_handler))

  predictions | "WriteOutputToGCS" >> beam.io.WriteToText( # pylint: disable=expression-not-assigned
    known_args.output,
    shard_name_template='',
    append_trailing_newlines=True)

  result = pipeline.run()
  result.wait_until_finish()
  return result