text: import torch
  import torch.nn as nn
  from torch.nn import functional as F
  import torch.autograd as autograd
  from torch.cuda.amp import autocast
  import numpy as np
  
  from dassl.engine import TRAINER_REGISTRY, TrainerX, SimpleNet
  from dassl.metrics import compute_accuracy
  from dassl.utils import count_num_param
  from dassl.optim import build_optimizer, build_lr_scheduler
  
  
  @TRAINER_REGISTRY.register()
  class RSC_DB(TrainerX):
      """RSC implementation from DomainBed."""
  
      def __init__(self, cfg):
          super().__init__(cfg)
          self.drop_f = (1 - cfg.TRAINER.RSC_DB.F_DROP_FACTOR) * 100
          self.drop_b = (1 - cfg.TRAINER.RSC_DB.B_DROP_FACTOR) * 100
  
      def build_model(self):
          cfg = self.cfg
  
          print('Building F')
          self.F = SimpleNet(cfg, cfg.MODEL, 0)
          self.F.to(self.device)
          print('# params: {:,}'.format(count_num_param(self.F)))
          self.optim_F = build_optimizer(self.F, cfg.OPTIM)
          self.sched_F = build_lr_scheduler(self.optim_F, cfg.OPTIM)
          self.register_model('F', self.F, self.optim_F, self.sched_F)
          fdim = self.F.fdim
  
          print('Building classifier')
          self.C = nn.Linear(fdim, self.num_classes)
          self.C.to(self.device)
          print('# params: {:,}'.format(count_num_param(self.C)))
          self.optim_C = build_optimizer(self.C, cfg.OPTIM)
          self.sched_C = build_lr_scheduler(self.optim_C, cfg.OPTIM)
          self.register_model('C', self.C, self.optim_C, self.sched_C)
  
      def forward_model_loss(self, all_x, all_y):
          # one-hot labels
          all_o = F.one_hot(all_y, self.num_classes)
          # features
          all_f = self.F(all_x)
          # predictions
          all_p = self.C(all_f)
  
          # Equation (1): compute gradients with respect to representation
          all_g = autograd.grad((all_p * all_o).sum(), all_f)[0]
  
          # Equation (2): compute top-gradient-percentile mask
          percentiles = np.percentile(all_g.cpu(), self.drop_f, axis=1)
          percentiles = torch.Tensor(percentiles)
          percentiles = percentiles.unsqueeze(1).repeat(1, all_g.size(1))
          mask_f = all_g.lt(percentiles.to(self.device)).float()
  
          # Equation (3): mute top-gradient-percentile activations
          all_f_muted = all_f * mask_f
  
          # Equation (4): compute muted predictions
          all_p_muted = self.C(all_f_muted)
  
          # Section 3.3: Batch Percentage
          all_s = F.softmax(all_p, dim=1)
          all_s_muted = F.softmax(all_p_muted, dim=1)
          changes = (all_s * all_o).sum(1) - (all_s_muted * all_o).sum(1)
          percentile = np.percentile(changes.detach().cpu(), self.drop_b)
          mask_b = changes.lt(percentile).float().view(-1, 1)
          mask = torch.logical_or(mask_f, mask_b).float()
  
          # Equations (3) and (4) again, this time mutting over examples
          all_p_muted_again = self.C(all_f * mask)
  
          # Equation (5): update
          loss = F.cross_entropy(all_p_muted_again, all_y)
  
          return loss, all_p_muted_again
  
      def forward_backward(self, batch):
          all_x, all_y = self.parse_batch_train(batch)
  
          if self.use_amp:
              with autocast():
                  loss, all_p_muted_again = self.forward_model_loss(all_x, all_y)
          else:
              loss, all_p_muted_again = self.forward_model_loss(all_x, all_y)
  
          self.model_backward_and_update(loss)
  
          loss_summary = {
              'loss': loss.item(),
              'acc': compute_accuracy(all_p_muted_again, all_y)[0].item()
          }
  
          if (self.batch_idx + 1) == self.num_batches:
              self.update_lr()
  
          return loss_summary
  
      def parse_batch_train(self, batch):
          input = batch['img']
          label = batch['label']
          input = input.to(self.device)
          label = label.to(self.device)
          return input, label
  
      def model_inference(self, input):
          return self.C(self.F(input))
  