code: def train_vision_model(model: SimpleConvNet, per_layer_scaling=False):
      loss_fn = nn.CrossEntropyLoss()
      optimizer = torch.optim.SGD(
          model.parameters(), lr=0.001, momentum=0.9
      )
  
      if torch.cuda.is_available():
          model.cuda()
  
      train_ds_loader = load_data("vision_model")
      model.train()
  
      layer_scale_dict = (
          {
              "conv1": 128,
              "conv2": 256,
              "fc1": 512,
              "fc2": 1024,
              "fc3": 8192,
          }
          if per_layer_scaling
          else {}
      )
      layer_scaler = LayerwiseGradientScaler(model, layer_scale_dict)
  
      for _ in range(2):
          for img, lbl in train_ds_loader:
              if torch.cuda.is_available():
                  img = img.cuda()
                  lbl = lbl.cuda()
  
              optimizer.zero_grad()
              layer_scaler.scale()
  
              predict = model(img)
              loss = loss_fn(predict, lbl)
  
              loss.backward()
  
              layer_scaler.unscale()
              layer_scaler.step(optimizer)
      return model
  
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 6 for calling a global function with example usage: train_vision_model(model, per_layer_scaling) and returns: model