I am creating a network for deep learning classification.
When I ran my code on the GPU, it raised this error with wrong size mismatch. (target size (5, 50, 76, 80) and torch.size([5]))
I am using MRI datasets (depth, height, width) with two labels 0 and 1 (no masks).
Here’s the code I used for creating dataset:
class BrainS18Dataset(Dataset):

    def __init__(self, img_list, sets): #delete root_dir
        with open(img_list, 'r') as f:
            self.img_list = [line.strip() for line in f]
            self.img_list.pop(0)
        self.input_D = sets.input_D
        self.input_H = sets.input_H
        self.input_W = sets.input_W
        self.phase = sets.phase

    def __onehot__(self, label):
        if label == '1':
            label = np.array([0, 1], dtype='f')
        elif label == '0':
            label = np.array([1, 0], dtype='f')
        
        return label
    
    def __nii2tensorarray__(self, data):
        [z, y, x] = data.shape
        new_data = np.reshape(data, [1, z, y, x])
        new_data = new_data.astype("float32")
            
        return new_data
    
    def __len__(self):
        return len(self.img_list)

    def __getitem__(self, idx):
        
        if self.phase == "train":
            # read image and labels
            ith_info = self.img_list[idx].split(",")

            
            img_name = os.path.join(ith_info[0])
            label = ith_info[1]
            img = nibabel.load(img_name)  # We have transposed the data from WHD format to DHW
            assert img is not None
            assert label is not None
            
            # data processing

            img_array = self.__training_data_process__(img)       
            label = self.__onehot__(label)
            
            # 2 tensor array
            img_array = self.__nii2tensorarray__(img_array)

            return img_array, label
        
        elif self.phase == "test":
            # read image
            ith_info = self.img_list[idx].split(",")
            
            img_name = os.path.join(ith_info[0])
            label = ith_info[1]
            

            img = nibabel.load(img_name)
            assert img is not None

            # data processing
            img_array = self.__training_data_process__(img)    

            # 2 tensor array
            img_array = self.__nii2tensorarray__(img_array)
            
            label = self.__onehot__(label)

            return img_array, label

    def __random_crop_to_same_size__(self, data): 
        
        """
        Crop to be the same size
        """

        [img_d, img_h, img_w] = data.shape
        data_processed = np.zeros((200, 300, 320))
        if [img_d, img_h, img_w] == [208, 300, 320]:    # not return label here
            data_processed = data[4:204, :, :]
        elif [img_d, img_h, img_w] == [226,320,320]:
            data_processed = data[13:213, 10:310, :]
        
        return data_processed

    def __training_data_process__(self, data): 

        data = data.get_fdata()
        data = self.__random_crop_to_same_size__(data)

        return data

And the training codes:
def train(data_loader, model, optimizer, scheduler, total_epochs, save_interval, save_folder, sets):
    # settings
    batches_per_epoch = len(data_loader)
    log.info('{} epochs in total, {} batches per epoch'.format(total_epochs, batches_per_epoch))

    loss_f = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([1.4]))

    print("Current setting is:")
    print(sets)
    print("\n\n")     
    if not sets.no_cuda:
        loss_f = loss_f.cuda()
    
    losses = []
    loss_plot = []
    
    #model.train()
    train_time_sp = time.time()
    for epoch in range(total_epochs):
        log.info('Start epoch {}'.format(epoch+1))
        
        val_losses = []
        
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()
                
            elif phase == 'val':
                print('Starting Validation: ')
                model.eval()
            z=0
        
            for batch_id, batch_data in enumerate(data_loaders[phase]):
                # getting data batch 
                z+=1
                batch_id_sp = epoch * batches_per_epoch
                volumes, labels = batch_data

                volumes = Variable(volumes.view(5,200,300,320))
                if not sets.no_cuda: 
                    volumes = volumes.cuda()

                if phase == 'train':
                    optimizer.zero_grad() 
                
                
                
                output = model(volumes)

                labels = Variable(torch.argmax(labels, dim=1))
                labels = labels.cuda()
                print(output.shape)
                print(labels.shape)
        
                loss = loss_f(output, labels)
                
                if phase == 'train':
                    losses.append(loss.item())
                    loss.backward()
                    optimizer.step()
                    
                    avg_batch_time = (time.time() - train_time_sp) / (1 + batch_id_sp)
                    
                    if False:
                        log.info(
                                '{} ||  Batch: {}-{} ({}), loss = {:.3f}, avg_batch_time = {:.3f}'\
                                .format(z, epoch, batch_id, batch_id_sp, loss.item(), avg_batch_time))
                  
                    if True:
                        # save model
                        if batch_id ==batches_per_epoch-1 and epoch % 1  == 0:
                            model_save_path = '{}/fdg_trail1_{}.pth.tar'.format(save_folder, epoch)
                            model_save_dir = os.path.dirname(model_save_path)
                            if not os.path.exists(model_save_dir):
                                os.makedirs(model_save_dir)
  
                            log.info('Save checkpoints: epoch = {}, batch_id = {}'.format(epoch, batch_id)) 
                            torch.save({
                                        'epoch': epoch,
                                        'batch_id': batch_id,
                                        'state_dict': model.state_dict(),
                                        'optimizer': optimizer.state_dict()},
                                        model_save_path)
   
                    loss_plot.append(np.mean(losses))
                    losses = []
                if phase == 'val':
                    val_losses.append(loss.item())
        if True:
            val_loss  = sum(val_losses)/len(val_losses)
            print('Epoch: {} || Validation Loss = {:.3f}'.format(epoch, val_loss))
            

    print('Finished training')      
                


if __name__ == '__main__':
    # settting
    sets = parse_opts()   
    
    # getting model
    
    torch.manual_seed(sets.manual_seed)
    model, parameters = generate_model(sets) 
    print (model)

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-5, weight_decay=5e-4)   
    scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.99)
    
    # train from resume
    if sets.resume_path:
        if os.path.isfile(sets.resume_path):
            print("=> loading checkpoint '{}'".format(sets.resume_path))
            checkpoint = torch.load(sets.resume_path)
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            print("=> loaded checkpoint '{}' (epoch {})"
              .format(sets.resume_path, checkpoint['epoch']))

    # getting data
    sets.phase = 'train'
    if sets.no_cuda:
        sets.pin_memory = False
    else:
        sets.pin_memory = True    
    training_dataset = BrainS18Dataset(sets.img_list, sets)
    validation_dataset = BrainS18Dataset(sets.val_list, sets)
    data_loader = DataLoader(training_dataset, batch_size=sets.batch_size, shuffle=True, num_workers=sets.num_workers, pin_memory=sets.pin_memory)     
    validation_loader = DataLoader(validation_dataset, batch_size=sets.batch_size, shuffle=True, num_workers=sets.num_workers, pin_memory=sets.pin_memory)    
    
    
    data_loaders = {"train": data_loader, "val": validation_loader}
    # training
    train(data_loaders, model, optimizer, scheduler, total_epochs=sets.n_epochs, save_interval=sets.save_intervals, save_folder=sets.save_folder, sets=sets) 


The error shows when using loss function (BCEWithLogitsLoss).
Also, for 2D image datasets, output = output.view(output, -1) is normally used to flatten the tensor in order to use the particular loss function, what is the common method for 3D images (with depths), can I also use view for the same circumstances?
Thanks for the help!