import torch.nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(3, 10, 3)
        self.conv2 = nn.Conv2d(10, 20, 3)
        self.conv3 = nn.Conv2d(20, 30, 3)
        self.maxpool = nn.MaxPool2d((2,2))
        self.linear = nn.Linear(120, 10)


    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.maxpool(x)
        
        x = F.relu(self.conv2(x))
        x = self.maxpool(x)
        
        x = F.relu(self.conv3(x))
        x = self.maxpool(x)
        
        x = x.view(x.shape[0], -1)
        x = self.linear(x)
        return x


network = Model() #example

res = []
for i in range(8): #should be: for inputs, labels in train_loader:
    inputs = torch.rand(32, 3, 32, 32) #example
    prototypes = network(inputs) # network: Resnet + FC layer.
    prototypes = prototypes.permute(1, 0)
    res.append(prototypes)
    
res = torch.cat(res, dim=1)
print(res.shape)
print(res.mean(dim=1).shape)

Here the mean is calculated over all outputs from the network - not the highest one. In the ideal situation - where there are no biases and the dataset is well balanced - all values should be similar.