Hi,
I am collecting the activations after the hidden layers.
I would like to do it only in the training forward pass and not in the validation forward pass.
This is my forward function:
def get_activations(self, l, x):
        if self.track:
            self.activations[l]['mean'].append(float(x.mean()))
            self.activations[l]['var'].append(float(x.mean()))

def forward(self, x):
        x = self.sig(self.fcI(x))
        for l in range(self.n_lay):
            x = self.sig(self.fcH[l](x))
            self.get_activations(l, x)
        x = self.fcO(x)
        return x

How could I specify get_activations to be run only if the network in in train mode?
@timeit
def train(model, criterion, optimizer, results, EPOCHS, LR):
    for epoch in range(EPOCHS):    
        
        # Training
        model.train()
        train_epoch(model, tr_loader, criterion, optimizer, LR, results)  ### Do it here
        
        # Validation
        model.eval()
        valid_epoch(model, ts_loader, criterion, results)          ### Do not do it here

Thank you!