Hi, thanks for your reply. I’m changing the number of output classes for the last linear layer in the line, model.last_linear.out_features = data.c. printing learn.model shows that the last layer has an output size of 37. So I’ve just taken the standard alexnet head, changed the number of outputs for the last layer and then passed that as a custom head. I made a new post which shows the printout information.
Edit: It turns out that the issue was related to what you pointed out. I thought I was editing the output of the last linear layer, but the way I did it was apparently not a proper way of doing so, and therefore the number of outputs when running learn.predict was actually 1000 even though learn.model said the output size was 37. So whereas I thought learn.summary was wrong, it was actually learn.model that was wrong.