def lr_schedule_fn()

in weak_to_strong/train.py [0:0]

• 5 lines of code
• 2 McCabe index (conditional complexity)

    def lr_schedule_fn(step):
        if lr_schedule == "constant":
            return 1
        else:
            assert False, f"invalid lr schedule, {lr_schedule}, must be constant or cosine_anneal"