I have a very large tensor of shape (512,3,224,224). I input it to model in batches of 32 and I then save the scores corresponding to the target label which is 2. in each iteration, after every slice, the shape of scores changes. Which leads to the following error. What am I doing wrong and how to fix it.
label = torch.ones(1)*2
def sub_forward(self, x):
    x = self.vgg16(x)
    x = self.bn1(x)
    x = self.linear1(x)
    x = self.linear2(x)
    return x


def get_scores(self, imgs, targets):
    b, _, _, _ = imgs.shape
    batch_size = 32
    total_scores = []
    for i in range(0, b, batch_size):
      scores = self.sub_forward(imgs[i:i+batch_size,:,:,:])
      scores = F.softmax(scores)
      labels = targets[i:i+batch_size]
      labels = labels.long()
      scores = scores[:,labels]
      print(i," scores: ", scores)
      total_scores.append(scores)
      print(i," total_socres: ", total_scores)
    total_scores = torch.stack(total_scores)
    return scores

0  scores:  tensor([[0.0811],
        [0.0918],
        [0.0716],
        [0.1680],
        [0.1689],
        [0.1319],
        [0.1556],
        [0.2966],
        [0.0913],
        [0.1238],
        [0.1480],
        [0.1215],
        [0.2524],
        [0.1283],
        [0.1603],
        [0.1282],
        [0.2668],
        [0.1146],
        [0.2043],
        [0.2475],
        [0.0865],
        [0.1869],
        [0.0860],
        [0.1979],
        [0.1677],
        [0.1983],
        [0.2623],
        [0.1975],
        [0.1894],
        [0.3299],
        [0.1970],
        [0.1094]], device='cuda:0')
0  total_socres:  [tensor([[0.0811],
        [0.0918],
        [0.0716],
        [0.1680],
        [0.1689],
        [0.1319],
        [0.1556],
        [0.2966],
        [0.0913],
        [0.1238],
        [0.1480],
        [0.1215],
        [0.2524],
        [0.1283],
        [0.1603],
        [0.1282],
        [0.2668],
        [0.1146],
        [0.2043],
        [0.2475],
        [0.0865],
        [0.1869],
        [0.0860],
        [0.1979],
        [0.1677],
        [0.1983],
        [0.2623],
        [0.1975],
        [0.1894],
        [0.3299],
        [0.1970],
        [0.1094]], device='cuda:0')]
32  scores:  tensor([], device='cuda:0', size=(32, 0))
32  total_socres:  [tensor([[0.0811],
        [0.0918],
        [0.0716],
        [0.1680],
        [0.1689],
        [0.1319],
        [0.1556],
        [0.2966],
        [0.0913],
        [0.1238],
        [0.1480],
        [0.1215],
        [0.2524],
        [0.1283],
        [0.1603],
        [0.1282],
        [0.2668],
        [0.1146],
        [0.2043],
        [0.2475],
        [0.0865],
        [0.1869],
        [0.0860],
        [0.1979],
        [0.1677],
        [0.1983],
        [0.2623],
        [0.1975],
        [0.1894],
        [0.3299],
        [0.1970],
        [0.1094]], device='cuda:0'), tensor([], device='cuda:0', size=(32, 0))]

`

RuntimeError: stack expects each tensor to be equal size, but got [32, 1] at entry 0 and [32, 0] at entry 1

`