code: def set_grad(
      model: nn.Module, torch_model: nn.Module, g_dtype: torch.dtype
  ) -> None:
      for p, torch_p in zip(
          model.parameters(), torch_model.parameters()
      ):
          torch_p.grad = torch.rand_like(torch_p)
          # avoid inconsistent grad and param dtype error
          orig_p = p.data
          p.data = torch_p.grad.clone().to(g_dtype)
          p.grad = p.data
          p.data = orig_p
  
file_path: ColossalAI/tests/test_optimizer/test_adam_optim.py
imports: from copy import deepcopy
  from typing import Type, Union
  import pytest
  import torch
  import torch.nn
  from torch.optim import Adam, AdamW
  from colossalai.nn.optimizer import CPUAdam, FusedAdam, HybridAdam
  from tests.kit.model_zoo import model_zoo
label_desc: Use image node_id 2 for calling a global function with example usage: set_grad(model, torch_model, g_dtype) without return types