Solved by dazzle-me in post #2 
              
              
                I guess it is doable. 
Consider an example : 
class MyModel(torch.nn.Module):
    def __init__(self):
         super(MyModel, self).__init__()
         self.current_epoch = 0

    def update_epoch(self, epoch):
         self.current_epoch = epoch

    def forward(self, x):
         if some_condition…