How to handle transforms.FiveCrop change in tensor size
I am trying to add transforms.FiveCrop to my model. I understand that this method of data augmentation adds dimensions to my tensor but I am not sure how to handle it. The documentation notes:

This transform returns a tuple of images and there may be a mismatch in the number of inputs and targets your Dataset returns. See below for an example of how to deal with this.

Example
&gt;&gt;&gt; transform = Compose([
&gt;&gt;&gt;    FiveCrop(size), # this is a list of PIL Images
&gt;&gt;&gt;    Lambda(lambda crops: torch.stack([ToTensor()(crop) for crop in crops])) # returns a 4D tensor
&gt;&gt;&gt; ])
&gt;&gt;&gt; #In your test loop you can do the following:
&gt;&gt;&gt; input, target = batch # input is a 5d tensor, target is 2d
&gt;&gt;&gt; bs, ncrops, c, h, w = input.size()
&gt;&gt;&gt; result = model(input.view(-1, c, h, w)) # fuse batch size and ncrops
&gt;&gt;&gt; result_avg = result.view(bs, ncrops, -1).mean(1) # avg over crops

...But I am not sure how to implement this.
Train loop:
    for batch_idx, (data, target) in enumerate(train_loader):

        print(f'Data: {data.shape}, Target: {target.shape}')       
        # Before Fivecrop: Data: torch.Size([32, 3, 224, 224]), Target: torch.Size([32])
        # After Fivecrop: Data: torch.Size([32, 5, 3, 224, 224]), Target: torch.Size([32])

        indx_target = target.clone()
        data = data.to(train_config.device)
        target = target.to(train_config.device)
        
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()

Can someone help explain how this is implemented in my train loop and what I am not understanding?
Thanks

This makes sense now. Here is for anyone who needs assistance. The new code looks as such:
    for batch_idx, (data, target) in enumerate(train_loader):
        
        bs, ncrops, c, h, w = data.size()
        print(f'Data: {data.view(-1, c, h, w).shape}, Target: {target.shape}')

        indx_target = target.clone()
        data = data.to(train_config.device)
        target = target.to(train_config.device)
        
        optimizer.zero_grad()
        output = model(data.view(-1, c, h, w)) # fuse batch size and ncrops
        output_avg = output.view(bs, ncrops, -1).mean(1) # average the output over fivecrops
        loss = F.cross_entropy(output_avg, target)
        loss.backward()
        optimizer.step()
        
        batch_loss = np.append(batch_loss, [loss.item()])
        prob = F.softmax(output_avg, dim=1)
        pred = prob.data.max(dim=1)[1]  
        correct = pred.cpu().eq(indx_target).sum()
        accuracy = float(correct) / float(len(data))
        batch_acc = np.append(batch_acc, [accuracy])

