def METHOD_NAME(self, grad, args):
    r"""Update the variables to take a single optimization step. Flattens and unflattens
    the inputs to maintain nested iterables as the parameters of the optimization.
    Args:
        grad (tuple [array]): the gradient of the objective
            function at point :math:`x^{(t)}`: :math:`\nabla f(x^{(t)})`
        args (tuple): the current value of the variables :math:`x^{(t)}`
    Returns:
        list [array]: the new values :math:`x^{(t+1)}`
    """
    args_new = list(args)
    trained_index = 0
    for index, arg in enumerate(args):
        if getattr(arg, "requires_grad", False):
            args_new[index] = arg - self.stepsize * grad[trained_index]
            trained_index += 1
    return args_new