code: def get_params_with_grad(trained_model):
      result = []
      for module_name, layer in trained_model.named_modules():
          if module_name != "":
              for param_name, param in layer.named_parameters():
                  if hasattr(param, "grad"):
                      logging.debug(
                          "testing equality for %s.%s"
                          % (module_name, param_name)
                      )
                      result.append(param.grad)
      return result
  
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 3 for calling a global function with example usage: get_params_with_grad(trained_model) and returns: result