def METHOD_NAME(self):
    with self.assertRaises(AssertionError):
        self.lpot_model.get_gradient("fc.bias")
    shape = None
    for name, tensor in self.lpot_model._model.named_parameters():
        if name == "fc.bias":
            shape = tensor.shape
            tensor.grad = torch.randn(shape)
            break
    new_grad = torch.zeros(shape)
    self.lpot_model.update_gradient("fc.bias", new_grad)
    assert torch.equal(torch.tensor(self.lpot_model.get_gradient("fc.bias")), torch.zeros(shape))
    rand_input = torch.rand(100, 3, 256, 256).float()
    rand_input.grad = torch.ones_like(rand_input)
    assert torch.equal(torch.tensor(self.lpot_model.get_gradient(rand_input)), torch.ones_like(rand_input))