import torch


def get_transformer_scheduler(optimizer, d_model, n_warmup, last_step=-1):
	"""
	Learning rate scheduler as described in "Attention is what you need".
	"""
	def lr_lambda(current_step):
		current_step += 1
		return max(d_model**-0.5 * min(current_step**-0.5, current_step*n_warmup**-1.5), 1e-5)
	return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=last_step)