#! /usr/bin/python # -*- encoding: utf-8 -*- import torch def Scheduler(optimizer, test_interval, max_epoch, lr_decay, **kwargs): sche_fn = torch.optim.lr_scheduler.StepLR(optimizer, step_size=test_interval, gamma=lr_decay) lr_step = 'epoch' print('Initialised step LR scheduler') return sche_fn, lr_step