What is the last layer being predicted? Meet a problem when trying to train the model

[luckui]

This is my codes, I use MSELoss(), and target is the size of torch.Size([4, 1, 1, 144, 256]), but the last layer is like torch.Size([4, 1, 64, 144, 256]), 64 is the num of hidden layers though.

for epoch in range(num_epochs):
        running_loss = 0.0
        for i, inputs in enumerate(dataloader):
            inputs = inputs[0].to(device)
            target = inputs[:, -1:, :, :, :]
            inputs = inputs[:, :-1, :, :, :]
            print(target.shape)
            optimizer.zero_grad()
            layerlist, last_states = model(inputs)
            # print(inputs.shape)
            # print(output[0].shape)
            h = last_states[0][0]
            c = last_states[0][1]
            layer = layerlist[-1]
            print(layer.shape)
            loss = criterion(layer, target)
            loss.backward()
            optimizer.step()
            running_loss+=loss.item()
        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(dataloader):.4f}")