def _check_one_layer_with_criterion(self, layer, criterion, *args, **kwargs):
    self.validator.validate(layer)
    for (name, param) in layer.named_parameters():
        if (('weight' in name) or ('bias' in name)):
            nn.init.uniform_(param, (- 1.0), 1.0)
    self._run_once(layer, criterion, *args)
    vanilla_run_grads = [(name, p.grad.detach()) for (name, p) in layer.named_parameters() if p.requires_grad]
    clipper = PerSampleGradientClipper(layer, ConstantFlatClipper(1000000000.0), batch_first=kwargs.get('batch_first', True), loss_reduction=criterion.reduction)
    self._run_once(layer, criterion, *args)
    for (param_name, param) in layer.named_parameters():
        if param.requires_grad:
            self.assertTrue(hasattr(param, 'grad_sample'), f"Per-sample gradients haven't been computed for {param_name}")
    clipper.clip_and_accumulate()
    clipper.pre_step()
    private_run_grads = [(name, p.grad.detach()) for (name, p) in layer.named_parameters() if p.requires_grad]
    for ((vanilla_name, vanilla_grad), (private_name, private_grad)) in zip(vanilla_run_grads, private_run_grads):
        assert (vanilla_name == private_name)
        self.assertTrue(torch.allclose(vanilla_grad, private_grad, atol=0.0001, rtol=0.01), f'Gradient mismatch. Parameter: {layer}.{vanilla_name}, loss: {criterion.reduction}')
    clipper.close()
def _check_one_layer_with_criterion(self, layer, criterion, *args, **kwargs):
    self.validator.validate(layer)
    for (name, param) in layer.named_parameters():
        if (('weight' in name) or ('bias' in name)):
            nn.init.uniform_(param, (- 1.0), 1.0)
    self._run_once(layer, criterion, *args)
    vanilla_run_grads = [(name, p.grad.detach()) for (name, p) in layer.named_parameters() if p.requires_grad]
    clipper = PerSampleGradientClipper(layer, ConstantFlatClipper(1000000000.0), batch_first=kwargs.get('batch_first', True), loss_reduction=criterion.reduction)
    self._run_once(layer, criterion, *args)
    for (param_name, param) in layer.named_parameters():
        if param.requires_grad:
            self.assertTrue(hasattr(param, 'grad_sample'), f"Per-sample gradients haven't been computed for {param_name}")
    clipper.clip_and_accumulate()
    clipper.pre_step()
    private_run_grads = [(name, p.grad.detach()) for (name, p) in layer.named_parameters() if p.requires_grad]
    for ((vanilla_name, vanilla_grad), (private_name, private_grad)) in zip(vanilla_run_grads, private_run_grads):
        assert (vanilla_name == private_name)
        self.assertTrue(torch.allclose(vanilla_grad, private_grad, atol=0.0001, rtol=0.01), f'Gradient mismatch. Parameter: {layer}.{vanilla_name}, loss: {criterion.reduction}')
    clipper.close()