def _get_gradients(self):
    gradients = []
    for p in self.model.parameters():
        if (p.grad is None):
            continue
        gradients.append(p.grad.flatten())
    gradients = torch.cat(gradients, dim=0)
    return gradients
def _get_gradients(self):
    gradients = []
    for p in self.model.parameters():
        if (p.grad is None):
            continue
        gradients.append(p.grad.flatten())
    gradients = torch.cat(gradients, dim=0)
    return gradients