code: def test_amp(use_amp: bool, fp16_recurrence: bool):
      its = 20
      warmup = 3
      batch_size = 200
      seq_len = 304
      hidden_size = 304
  
      input_size = hidden_size
      device = "cuda"
  
      inputs = torch.rand(
          seq_len, batch_size, input_size, device=device
      )
      cell = sru_lib.SRUCell(
          input_size, hidden_size, amp_recurrence_fp16=fp16_recurrence
      ).to(device)
      opt = optim.Adam(lr=0.02, params=cell.parameters())
  
      if use_amp:
          amp: Any = torch.cuda.amp
      else:
          amp = NullAmp()
      scaler = amp.GradScaler()
  
      sync(device)
      for it in range(its + warmup):
          if it == warmup:
              sync(device)
              start_time = time.time()
          with amp.autocast():
              out = cell(inputs)
              s = out[0].sum()
  
          opt.zero_grad()
          scaler.scale(s).backward()
          scaler.step(opt)
          scaler.update()
      sync(device)
      elapsed = time.time() - start_time
      print("elapsed %.3f" % elapsed)
  
file_path: sru/test/test_amp.py
imports: import pytest
  import sru
  from typing import Any
  import time
  import torch
  from torch import optim
  import contextlib
label_desc: Use image node_id 2 for calling a global function with example usage: test_amp(use_amp, fp16_recurrence) without return types