How to implement an epoch-dependent parameter inside a neural network in Pytorch?
For example, let's say I want to use Softmax with temperature after a layer and I also want to decrease the temperature after 5 epochs.
for epoch in range(EPOCHS):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
    inputs, labels = data[0].to(torch_device), data[1].to(torch_device)
    optimizer.zero_grad()
    outputs = my_model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    my_model.adjust_softmax_temperature(epoch)
    # ...

How should I implement the adjust_softmax_temperature method in order to make my code to work?


as far as I know, the model is loaded BEFORE thre training, and the forward step passes data through the already-defined model. Am I missing something here?

You can mutate the state of your model even after it has been defined. In this case, we're interested in changing a property of my_model during training.
Here is a broad overview of how you could implement this.
First, expose a function to update the model's temperature based on the provided epoch number:
class MyModel(nn.Module):
    def __init__(self, t0):
        self.t = t0
    
    def forward(self, x):
        # compute forward pass using self.t
        pass

    def adjust_softmax_temperature(self, epoch):
        if epoch % 5:
            self.t /= 10 # update policy of temperature

Then, after having defined your model, you can call the update function at the end of each epoch: it will decide based on epoch wether to update self.t or not.
my_model = MyModel(t0=100)

for epoch in range(EPOCHS):
    for i, data in enumerate(trainloader, 0):
        # train iteration
    my_model.adjust_softmax_temperature(epoch)

