RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 512 but got size 100 for tensor number 1 in the list

I am trying to use a slightly modified GAN model.
When each data was input one by one without using batch, the generator worked normally.
However, when inputting data into generator using batch_size through dataloader, an error such as the title occurs.
I’m curious which part is the problem.

class add_Linear(nn.Module): Could this be the problem?

The size of data before generator input is,
When entering one by one,
condition_im : 3,256,256
radomixxxx: 100

When batchsize is 10,
condition_im :10, 3,256,256
radomixxxx: 10,100

Thanks for your help.

# version 02
# basic model + Radiomix data(concat 1x100) + augmentation
#
import torch

# GPU 사용 가능 -> True, GPU 사용 불가 -> False
print(torch.cuda.is_available())

import os
import glob
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

from torchsummary import summary

import torch
import torch.nn as nn

from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.utils import save_image

import scipy.io as sio
import mat73

condition_data_im = sio.loadmat('data/second/f_condition_data_im_256.mat')
condition_data_im = condition_data_im['f_condition_data_im']
condition_data_im = np.expand_dims(condition_data_im,0)

condition_data_seg = sio.loadmat('data/second/f_condition_data_seg_256.mat')
condition_data_seg = condition_data_seg['f_condition_data_seg']
condition_data_seg = np.expand_dims(condition_data_seg,0)

condition_data_RD = sio.loadmat('data/second/f_condition_data_RD_256.mat')
condition_data_RD = condition_data_RD['f_condition_data_RD']
condition_data_RD = np.expand_dims(condition_data_RD,0)

condition_data_all = np.concatenate((condition_data_im,condition_data_seg,condition_data_RD),axis=0)

real_data = sio.loadmat('data/second/f_real_data_im_256.mat')
real_data = real_data['f_real_data_im']
real_data = np.expand_dims(real_data,0)


radiomix_data = np.zeros([100,695])




class ImageDataset(Dataset):
    def __init__(self,condi ,real, mix, transforms_=None):
        self.transform = transforms_
        self.condi = torch.FloatTensor(condi).permute(3,0,1,2) #(4,0,1,2,3) #
        self.real = torch.FloatTensor(real).permute(3,0,1,2)#(4,0,1,2,3)
        self.mix = torch.FloatTensor(mix).permute(1,0)
        self.len = self.condi.shape[0]

    def __getitem__(self, idx):

         return self.condi[idx], self.real[idx], self.mix[idx]

    def __len__(self):
        return self.len

#np.size(condition_data_all,4)

transforms_ = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train_dataset = ImageDataset(condition_data_all, real_data, radiomix_data, transforms_=None)

train_dataloader = DataLoader(train_dataset, batch_size=10, shuffle=True)

#train_dataset[0][0].size()
train_features = next(iter(train_dataloader))
qwe =  train_features[0]
ewq =  train_features[2]




# U-Net 아키텍처의 다운 샘플링(Down Sampling) 모듈
class UNetDown(nn.Module):
    def __init__(self, in_channels, out_channels, normalize=True, dropout=0.0):
        super(UNetDown, self).__init__()
        # 너비와 높이가 2배씩 감소
        layers = [nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False)]
        if normalize:
            layers.append(nn.InstanceNorm2d(out_channels))
        layers.append(nn.LeakyReLU(0.2))
        if dropout:
            layers.append(nn.Dropout(dropout))
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)


# U-Net 아키텍처의 업 샘플링(Up Sampling) 모듈: Skip Connection 사용
class UNetUp(nn.Module):
    def __init__(self, in_channels, out_channels, dropout=0.0):
        super(UNetUp, self).__init__()
        # 너비와 높이가 2배씩 증가
        layers = [nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False)]
        layers.append(nn.InstanceNorm2d(out_channels))
        layers.append(nn.ReLU(inplace=True))
        if dropout:
            layers.append(nn.Dropout(dropout))
        self.model = nn.Sequential(*layers)

    def forward(self, x, skip_input):
        x = self.model(x)
        x = torch.cat((x, skip_input), 1) # 채널 레벨에서 합치기(concatenation)

        return x

class add_Linear(nn.Module):
    def __init__(self, in_channels, out_channels=512):
        super(add_Linear, self).__init__()

        layers = [nn.Linear(in_channels, out_channels)]
        self.model = nn.Sequential(*layers)

    def forward(self, x, radiomix):
        sq_x = torch.squeeze(x)
        sq_x2 = torch.squeeze(sq_x)
        cat_rad = torch.cat((sq_x2, radiomix))
        x = self.model(cat_rad)
        ex_cat2 =  torch.unsqueeze(x,1) #
        ex2_cat2 = torch.unsqueeze(ex_cat2, 1)  #
        ex3_cat2 = torch.unsqueeze(ex2_cat2, 0)  #

        return ex3_cat2


# U-Net 생성자(Generator) 아키텍처
class GeneratorUNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1):
        super(GeneratorUNet, self).__init__()

        self.down1 = UNetDown(in_channels, 64, normalize=False) # 출력: [64 X 128 X 128]
        self.down2 = UNetDown(64, 128) # 출력: [128 X 64 X 64]
        self.down3 = UNetDown(128, 256) # 출력: [256 X 32 X 32]
        self.down4 = UNetDown(256, 512, dropout=0.5) # 출력: [512 X 16 X 16]
        self.down5 = UNetDown(512, 512, dropout=0.5) # 출력: [512 X 8 X 8]
        self.down6 = UNetDown(512, 512, dropout=0.5) # 출력: [512 X 4 X 4]
        self.down7 = UNetDown(512, 512, dropout=0.5) # 출력: [512 X 2 X 2]
        self.down8 = UNetDown(512, 512, normalize=False, dropout=0.5) # 출력: [512 X 1 X 1]
        self.ADD = add_Linear(612, 512)  # [512]

        # Skip Connection 사용(출력 채널의 크기 X 2 == 다음 입력 채널의 크기)
        self.up1 = UNetUp(512, 512, dropout=0.5) # 출력: [1024 X 2 X 2]
        self.up2 = UNetUp(1024, 512, dropout=0.5) # 출력: [1024 X 4 X 4]
        self.up3 = UNetUp(1024, 512, dropout=0.5) # 출력: [1024 X 8 X 8]
        self.up4 = UNetUp(1024, 512, dropout=0.5) # 출력: [1024 X 16 X 16]
        self.up5 = UNetUp(1024, 256) # 출력: [512 X 32 X 32]
        self.up6 = UNetUp(512, 128) # 출력: [256 X 64 X 64]
        self.up7 = UNetUp(256, 64) # 출력: [128 X 128 X 128]

        self.final = nn.Sequential(
            nn.Upsample(scale_factor=2), # 출력: [128 X 256 X 256]
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(128, out_channels, kernel_size=4, padding=1), # 출력: [3 X 256 X 256]
            #nn.Tanh(),
        )

    def forward(self, x, radiomix):
        # 인코더부터 디코더까지 순전파하는 U-Net 생성자(Generator)
        d1 = self.down1(x)
        d2 = self.down2(d1)
        d3 = self.down3(d2)
        d4 = self.down4(d3)
        d5 = self.down5(d4)
        d6 = self.down6(d5)
        d7 = self.down7(d6)
        d8 = self.down8(d7)
        add_d8 = self.ADD(d8, radiomix)
        u1 = self.up1(add_d8, d7)
        u2 = self.up2(u1, d6)
        u3 = self.up3(u2, d5)
        u4 = self.up4(u3, d4)
        u5 = self.up5(u4, d3)
        u6 = self.up6(u5, d2)
        u7 = self.up7(u6, d1)

        return self.final(u7)

generator = GeneratorUNet()
generator.cuda()

#summary(generator, input_size=(3,256,256))



# U-Net 판별자(Discriminator) 아키텍처
class Discriminator(nn.Module):
    def __init__(self, in_channels=2):
        super(Discriminator, self).__init__()

        def discriminator_block(in_channels, out_channels, normalization=True):
            # 너비와 높이가 2배씩 감소
            layers = [nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1)]
            if normalization:
                layers.append(nn.InstanceNorm2d(out_channels))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.model = nn.Sequential(
            # 두 개의 이미지(실제/변환된 이미지, 조건 이미지)를 입력 받으므로 입력 채널의 크기는 2배
            *discriminator_block(in_channels * 2, 64, normalization=False), # 출력: [64 X 128 X 128]
            *discriminator_block(64, 128), # 출력: [128 X 64 X 64]
            *discriminator_block(128, 256), # 출력: [256 X 32 X 32]
            *discriminator_block(256, 512), # 출력: [512 X 16 X 16]
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(512, 1, kernel_size=4, padding=1, bias=False) # 출력: [1 X 16 X 16]
        )

    # img_A: 실제/변환된 이미지, img_B: 조건(condition)
    def forward(self, img_A, img_B):
        # 이미지 두 개를 채널 레벨에서 연결하여(concatenate) 입력 데이터 생성
        img_input = torch.cat((img_A, img_B), 1)
        return self.model(img_input)


def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find("BatchNorm2d") != -1:
        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
        torch.nn.init.constant_(m.bias.data, 0.0)


# 생성자(generator)와 판별자(discriminator) 초기화
discriminator = Discriminator()
discriminator.cuda()


# 가중치(weights) 초기화
generator.apply(weights_init_normal)
discriminator.apply(weights_init_normal)

# 손실 함수(loss function)
criterion_GAN = torch.nn.MSELoss()
criterion_pixelwise = torch.nn.L1Loss()

criterion_GAN.cuda()
criterion_pixelwise.cuda()

# 학습률(learning rate) 설정
lr = 0.0005 #0.0002

# 생성자와 판별자를 위한 최적화 함수
optimizer_G = torch.optim.Adam(generator.parameters(), lr=lr, betas=(0.5, 0.999))
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=lr, betas=(0.5, 0.999))


import time

n_epochs = 30 # 학습의 횟수(epoch) 설정

# 변환된 이미지와 정답 이미지 사이의 L1 픽셀 단위(pixel-wise) 손실 가중치(weight) 파라미터
lambda_pixel = 100

start_time = time.time()


for epoch in range(n_epochs):
    for i, batch in enumerate(train_dataloader):
        # 모델의 입력(input) 데이터 불러오기


        condition_im = batch[0].cuda()
        real_im = batch[1].cuda()
        radiomixxx = batch[2].cuda()

         # 진짜(real) 이미지와 가짜(fake) 이미지에 대한 정답 레이블 생성 (너바와 높이를 16씩 나눈 크기)
        real = torch.cuda.FloatTensor(real_im.size(0), 1, 16, 16).fill_(1.0) # 진짜(real): 1
        fake = torch.cuda.FloatTensor(real_im.size(0), 1, 16, 16).fill_(0.0) # 가짜(fake): 0

        """ 생성자(generator)를 학습합니다. """
        optimizer_G.zero_grad()

        # 이미지 생성
        fake_im = generator(x = condition_im, radiomix = radiomixxx) ##### @@@@ ERROR


Traceback (most recent call last):
  File "/home/smc-gpu/anaconda3/envs/CT_GAN/lib/python3.7/code.py", line 90, in runcode
    exec(code, self.locals)
  File "<input>", line 1, in <module>
  File "/home/smc-gpu/anaconda3/envs/CT_GAN/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "<input>", line 98, in forward
  File "/home/smc-gpu/anaconda3/envs/CT_GAN/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "<input>", line 48, in forward
RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 512 but got size 100 for tensor number 1 in the list.

I guess the error is raised in an invalid torch.cat operation as seen here:

sq_x2 = torch.randn(8, 512)
radiomix = torch.randn(8, 100)
cat_rad = torch.cat((sq_x2, radiomix))
# RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 512 but got size 100 for tensor number 1 in the list.

Are you sure you want to concatenate these tensors in dim0 and not in e.g. dim1?