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(self.current_epoch):
             ...
         else:
             ...

Before every epoch of your training, you can update epoch inside the model and that’s it 
for epoch in range(1, epochs + 1):
     model.update_epoch(epoch)
     for data, labels in dataloader:
         ...