Hello, I am attempting to recreate the CURE regularizer (https://arxiv.org/abs/1811.09716 6) and have a question about how to compute the gradient of the regularization term. Specifically, the regularizer is

and is used in conjunction with the standard cross entropy loss of a DNN classifier. I am following the implementation (from snippet) and may be misunderstanding the key line for the gradient computation. Here is how the repo computes the regularizer:

  
      github.com
  
  
    F-Salehi/CURE_robustness/blob/daeef24ea443c37a55f87f93679a44d8504d88e7/CURE/CURE.py#L200-L217

def regularizer(self, inputs, targets, h = 3., lambda_ = 4):
    '''
    Regularizer term in CURE
    '''
    z, norm_grad = self._find_z(inputs, targets, h)
    
    inputs.requires_grad_()
    outputs_pos = self.net.eval()(inputs + z)
    outputs_orig = self.net.eval()(inputs)


    loss_pos = self.criterion(outputs_pos, targets)
    loss_orig = self.criterion(outputs_orig, targets)
    grad_diff = torch.autograd.grad((loss_pos-loss_orig), inputs, grad_outputs=torch.ones(targets.size()).to(self.device),
                                    create_graph=True)[0]
    reg = grad_diff.reshape(grad_diff.size(0), -1).norm(dim=1)
    self.net.zero_grad()


    return torch.sum(self.lambda_ * reg) / float(inputs.size(0)), norm_grad



  
  
    
    
  
  


My question is, is this:
grad_diff = torch.autograd.grad((loss_pos-loss_orig), inputs, grad_outputs=torch.ones(targets.size()).to(self.device), create_graph=True)[0]

the correct statement for the gradient computation involved in the regularization term? At first glance, it appears to be allowing for the computation of the gradient of the difference of the losses, rather than the difference of the gradients. But, I may be misinterpreting this line completely.
Any comments on if this is the correct/incorrect way to compute the regularization term and potential solutions or improvements would be greatly appreciated.
Thanks,
Nate