I have just started using the Lightning framework, and my model is a CNN followed by a linear classification layer. Data is first passed through the CNN, and then from the classification layer. Before passing the data through the classification layer, I optionally use a dropout layer, and then a global max pooling. The data is sequences, thus batches do not have the same size, so the global max pooling is formed dynamically with each batch.
full_features_output = self.cnn(x)
pool_output = F.max_pool1d(full_features_output, kernel_size=full_features_output.shape[2])
These steps are repeated in train, evaluation and inference (dropout, is deactivated in evaluation and inference mode). Since Lightning organizes the code, I think that it is not the best practice to replicate this code in train_step(), eval_step() and test_step().
I found the LightningModule.forward() method which seems relevant, but the documentation says
To run data through your model only (separate from training_step)
and there were other posts that I read, mentioning also that this method is run in inference.
Should I use this method in this case? If not, what is the recommended way to structure the code?
Regards