Using a target size (torch.Size([64])) that is different to the input size (torch.Size([1])). This will likely lead to incorrect results

Question

input_size = 200000
hidden_size = 2
num_layers = 2
num_classes = 2
batch_size = 64
num_epochs=5

This is my simple model

class GRU(nn.Module):
def __init__(self,input_size,hidden_size,num_layers,num_classes,batch_size):
    super(GRU,self).__init__()
    self.num_layers=num_layers
    self.hidden_size=hidden_size
    self.input_size=input_size
    self.num_classes=num_classes
    
    self.GRU_0=nn.GRU(input_size,hidden_size,num_layers*2,batch_first=True)
    self.fc = nn.Linear(128, 1)
   
   
def forward(self,x):
    h0=torch.zeros(self.num_layers*2,batch_size,self.hidden_size)
    h0=weight_init(h0)
    x=x.type(torch.float32)
    #Forward Propagation
    out,_=self.GRU_0(x,h0)
    out=out.flatten()
    out=self.fc(out)
    
    return out

Training Loop

for epoch in tqdm(range(num_epochs)):
for j,data in enumerate(train_loader):
    inputs,labels=data
    outputs=model(inputs)
    print(outputs.shape)
    labels=labels.to(torch.float32)
    loss = criterion(outputs, labels)
    
    # Backward and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if (j+1) % 64 == 0:
        print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
               .format(epoch+1, num_epochs, i+1, len(train_loader), loss.item()))

I have to get output torch.Size([64,1]) because i set batch_size as 64. But instead of i get torch.Size([1]). Can Someone help me to resolve this problem.

Answer 1

You have a problem in your forward:

out=out.flatten()

Flats the entire output, including the batch dimension. Use start_dim:

out=out.flatten(start_dim=1)