this should work.
def _lr_lambda(current_step):
        """
        _lr_lambda returns a multiplicative factor given an interger parameter epochs.
        """

        if current_step < 400:
            _lr =.1
        elif current_step < 32000:
            _lr = 1
        elif current_step < 48000:
            _lr = .1
        else:
            _lr = .01

        return _lr

scheduler = LambdaLR(optimizer, _lr_lambda, last_epoch)