Model(input_data) only produces 10 output classes

model(input_data) or model.forward(input_data) is used to get the logits for all the classes.

My ‘model’ is a CNN and has final fc layer with out_features=200. ‘input_data’ is a tensor of size ([1, 160, 3, 224, 224]).

logits = model(input_data) should produce the logits for 200 classes, but ‘logits’ is of size ([10, 1]).

I don’t see where the number 10 is being set - I don’t use it anywhere in my code.

Any help fixing this would be greatly appreciated.

Thank you,

I solved it, the problem was with how the output probabilities were being averaged.