code: def test_custom_v(gpu=False):
      x = torch.randn(10, 2, 5)
      cell_1 = SRUCell(5, 5, bidirectional=True)
      if gpu:
          x = x.cuda()
          cell_1 = cell_1.cuda()
      weight = cell_1.weight
      bias = cell_1.bias
      weight_c = cell_1.weight_c
      h_1, c_1 = cell_1(x)
      loss_1 = h_1.sum()
      loss_1.backward()
      print("----------")
      print("SRU w/o custom_m:")
      print("loss: {}".format(loss_1))
      print("c: {}".format(c_1.sum()))
      print("grad w: {}".format(weight.grad.sum()))
      print("grad v: {}".format(weight_c.grad.sum()))
      print("")
  
      cell_1.zero_grad()
      weight_c_custom = (
          weight_c.view(2, -1).transpose(0, 1).contiguous().view(-1)
      )
  
      # weight_c is (2, bidir, d)
      # but custom weight_c is providing (length, batch, bidir, d, 2)
      def custom_m(input, **kwargs):
          U = input.matmul(weight)
          V = input.new_zeros(
              input.size(0), input.size(1), weight_c_custom.size(0)
          )
          return U, V
  
      cell_2 = SRUCell(5, 5, bidirectional=True, custom_m=custom_m)
      cell_2.weight_c.data.copy_(weight_c_custom)
      cell_2.bias = bias
      if gpu:
          cell_2 = cell_2.cuda()
      h_2, c_2 = cell_2(x)
      loss_2 = h_2.sum()
      loss_2.backward()
      print("SRU w/ custom_m:")
      print("loss: {}".format(loss_2))
      print("c: {}".format(c_2.sum()))
      print("grad w: {}".format(weight.grad.sum()))
      print("grad v: {}".format(cell_2.weight_c.grad.sum()))
      print("----------")
      print("")
  
imports: import cProfile, pstats, io
  import torch
  from sru import SRUCell
label_desc: Use image node_id 4 for calling a global function with example usage: test_custom_v(gpu) and returns: U,  V