code: def get_rand_tensors(global_shape, device):
      inp_t = torch.rand(
          global_shape, dtype=torch.float32, device=device
      ).to(memory_format=torch.channels_last)
      weight = torch.rand(
          global_shape[1], dtype=torch.float32, device=device
      )
      bias = torch.rand(
          global_shape[1], dtype=torch.float32, device=device
      )
      _grad_out = torch.rand(
          global_shape, dtype=torch.float32, device=device
      ).to(memory_format=torch.channels_last)
      return inp_t, weight, bias, _grad_out
  
imports: import copy
  import typing
  import unittest
  import torch
  import torch.nn
  from torch.testing._internal import common_utils
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
label_desc: Use image node_id 1 for calling a global function with example usage: get_rand_tensors(global_shape, device) and returns: inp_t,  weight,  bias,  _grad_out