code: def testGradScalerCapturable(self):
      params_ = [p for p in self.model_.parameters() if p.requires_grad]
      optimizer_ = apex.optimizers.FusedAdam(
          params_, lr=self.lr, capturable=True
      )
      scaler = torch.cuda.amp.GradScaler(enabled=True)
      scaler_ = torch.cuda.amp.GradScaler(enabled=True)
  
      for i in range(100):
          x = (
              torch.rand([32, 1, 28, 28])
              .cuda()
              .to(memory_format=torch.channels_last)
          )
          x_ = x.clone()
          gt = torch.rand([32, 10]).cuda()
          gt_ = gt.clone()
  
          # Reference
          with torch.cuda.amp.autocast(enabled=True):
              y = self.model(x)
              loss = ((gt - y) ** 2).mean()
  
          scaler.scale(loss).backward()
          scaler.step(self.optimizer)
          scaler.update()
  
          # DUT
          with torch.cuda.amp.autocast(enabled=True):
              y = self.model_(x)
              loss_ = ((gt_ - y) ** 2).mean()
  
          scaler_.scale(loss_).backward()
          scaler_.step(optimizer_)
          scaler_.update()
  
          for module in zip(
              self.model.modules(), self.model_.modules()
          ):
              m = module[0]
              m_ = module[1]
              if isinstance(m, nn.Conv2d) or isinstance(m_, nn.Linear):
                  torch.testing.assert_close(
                      m.weight,
                      m_.weight,
                      atol=1e-3,
                      rtol=1e-3,
                      equal_nan=True,
                  )
                  torch.testing.assert_close(
                      m.weight.grad,
                      m_.weight.grad,
                      atol=1e-3,
                      rtol=1e-3,
                      equal_nan=True,
                  )
  
          # Init for next iteration
          self.optimizer.zero_grad()
          optimizer_.zero_grad()
  
          self.model_.load_state_dict(
              copy.deepcopy(self.model.state_dict())
          )
  
file_path: apex/tests/L0/run_optimizers/test_adam.py
imports: import copy
  import math
  import random
  import unittest
  import torch
  import torch.nn.functional
  from torch import nn
label_desc: Use image node_id 3 for calling the AdamTest obj's underlying member method code with example usage: obj.testGradScalerCapturable() without return types