I was wondering if we could explore dynamically setting the out_dim parameter in the setup method of the LightningModule , leveraging the fact that it inherits DataHooks. This way, the model could fetch the number of labels (num_of_labels) directly from the data module after it's loaded.

By doing this, we could avoid requiring the out_dim parameter to be explicitly passed in the command line every time. Instead, the model could dynamically adapt based on the dataset being used. This would simplify the workflow and reduce potential user errors.

Originally posted by @aditya0by0 in #64 (comment)