Finding the Number of Parameters given an Input Shape


I wanted to write a variant of PReLU with one main distinction: it is always trainable for each channel-wise parameter instead of by layer.
The official PReLU implementation requires the user to enter the nchannels if it is not layer by layer.
Is there a way to discern that information given just the input tensor?

If I understand your use case correctly, you would like to use the first input to initialize the internal parameters of your custom activation module?
If so, take a look at .g. LazyLinear, which does the same and try to adapt its usage to your module.

1 Like

Thanks for the info- I will check out the Lazy Linear function and see if it can be used for my purposes.