code: def test_vision_model() -> None:
      # The os.environ below doesn't seem to be enough if the test is run on CI with many other tests
      # together.
      # see: https://app.circleci.com/pipelines/github/facebookresearch/fairscale/4086/workflows/72b1470a-55f8-4a45-afe5-04641b093bef/jobs/45179/tests#failed-test-0
      # Skipping for now.
      # Also, TODO (Min): improving downloading code above before re-enable this.
      skip_a_test_if_in_CI()
      # Remove randomness from various sources while testing.
      torch.use_deterministic_algorithms(True)  # type: ignore
      # set environment variable in CircleCI for test to pass: CUBLAS_WORKSPACE_CONFIG = :4096:8
      os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
  
      m1 = SimpleConvNet()
      m2 = SimpleConvNet()
  
      vision_model = train_vision_model(m1, False)
      scaled_vision_model = train_vision_model(m2, True)
  
      for elt in zip(
          get_params_with_grad(vision_model),
          get_params_with_grad(scaled_vision_model),
      ):
          assert torch.allclose(elt[0], elt[1])
  
imports: import logging
  import os
  from typing import Any, List, Tuple, Union
  import numpy
  import pytest
  from sklearn.datasets import make_blobs
  import torch
  from torch.cuda.amp.autocast_mode import autocast
  import torch.nn
  from torch.utils.data import DataLoader
  import torchvision
  import torchvision.transforms
  from fairscale.fair_dev.common_paths import DATASET_CACHE_ROOT
  from fairscale.fair_dev.testing.testing import skip_a_test_if_in_CI
  from fairscale.optim.layerwise_gradient_scaler import LayerwiseGradientScaler
label_desc: Use image node_id 7 for calling a global function with example usage: test_vision_model() without return types