def update_learning_rate(
        self, optimizer: "torch.optim.optimizer.Optimizer", epoch: int, nb_epochs: int, lr_schedule: str = "step"
    ) -> None:
        """
        adjust learning rate of the optimizer.

        :param optimizer: optimizer of the classifier.
        :param epoch: current training epoch.
        :param nb_epochs: total training epoch.
        :param lr_schedule: string denoting learning rate scheduler for optimizer
        :return: calculated learning rate
        """
        if lr_schedule.lower() == "cosine":
            if self._train_params["oaat_warmup"] == 1:
                l_r = self._train_params["lr"] * 0.5 * (1 + np.cos((epoch - 10) / (nb_epochs - 10) * np.pi))
            else:
                l_r = self._train_params["lr"] * 0.5 * (1 + np.cos(epoch / nb_epochs * np.pi))

            for param_group in optimizer.param_groups:
                param_group["lr"] = l_r

        elif lr_schedule.lower() == "linear":
            l_r = (epoch + 1) * (self._train_params["lr"] / 10)

            for param_group in optimizer.param_groups:
                param_group["lr"] = l_r

        elif lr_schedule.lower() == "step":
            if epoch >= 75 * nb_epochs / 110:
                l_r = self._train_params["lr"] * 0.1
            if epoch >= 90 * nb_epochs / 110:
                l_r = self._train_params["lr"] * 0.01
            if epoch >= 100 * nb_epochs / 110:
                l_r = self._train_params["lr"] * 0.001

            for param_group in optimizer.param_groups:
                param_group["lr"] = l_r
        else:
            raise ValueError(f"lr_schedule {lr_schedule} not supported")

        adjust learning rate of the optimizer.

        :param optimizer: optimizer of the classifier.
        :param epoch: current training epoch.
        :param nb_epochs: total training epoch.
        :param lr_schedule: string denoting learning rate scheduler for optimizer
        :return: calculated learning rate
        
art/defences/trainer/adversarial_trainer_oaat_pytorch.py