Source code for flambe.optim.scheduler

import torch
from flambe.compile import Component


[docs]class LRScheduler(torch.optim.lr_scheduler._LRScheduler, Component):

[docs]    def state_dict(self):
        state_dict = super().state_dict()
        del state_dict['_schema']
        del state_dict['_saved_kwargs']
        del state_dict['_extensions']
        return state_dict


[docs]class LambdaLR(torch.optim.lr_scheduler.LambdaLR, Component):

[docs]    def state_dict(self):
        state_dict = super().state_dict()
        del state_dict['_schema']
        del state_dict['_saved_kwargs']
        del state_dict['_extensions']
        return state_dict