Custom layer caused the batch dimension mismatch in pytorch. How to fix this problem?
I have tried to train a GCN model.I defined the custom layer I needed. However, It cause some dimension mismatch when I do some batch training.
the codes are as following :
import math
import numpy as np
import torch

import torch.nn as nn
from torch.nn.parameter import Parameter
from torch.nn.modules.module import Module
import torch.nn.functional as F
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from tqdm import tqdm

# =============================================================================
# model define
# =============================================================================
class GraphConvolution(Module):
    &quot;&quot;&quot;
    Simple GCN layer, similar to https://arxiv.org/abs/1609.02907
    &quot;&quot;&quot;

    def __init__(self, in_features, out_features, bias=True):
        super(GraphConvolution, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
        if bias:
            self.bias = Parameter(torch.FloatTensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

    def forward(self, input, adj):
        support = torch.mm(input, self.weight)
        output = torch.spmm(adj, support)
        if self.bias is not None:
            return output + self.bias
        else:
            return output

    def __repr__(self):
        return self.__class__.__name__ + ' (' \
               + str(self.in_features) + ' -&gt; ' \
               + str(self.out_features) + ')'
               
          
class GCN(nn.Module):
    def __init__(self, nfeat, nhid, nclass):
        super(GCN, self).__init__()

        self.gc1 = GraphConvolution(nfeat, nhid)
        self.gc2 = GraphConvolution(nhid, nclass)
        self.linear = nn.Linear(nclass, 1)
        # self.dropout = dropout

    def forward(self, x, adj):
        x = F.relu(self.gc1(x, adj))
        # x = F.dropout(x, self.dropout, training=self.training)
        x =  F.relu(self.gc2(x, adj))
        x = self.linear(x)
        return x



def train(dataloader, model, loss_fn, optimizer,adj):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred = model(X,adj)
        loss = loss_fn(pred, y)

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f&quot;loss: {loss:&gt;7f}  [{current:&gt;5d}/{size:&gt;5d}]&quot;)

def test(dataloader, model, loss_fn,adj):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X,adj)
            test_loss += loss_fn(pred, y).item()
            # correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num_batches
    # correct /= size
    # Accuracy: {(100*correct):&gt;0.1f}%,
    print(f&quot;Test Error: \n  Avg loss: {test_loss:&gt;8f} \n&quot;)
    
    

when I run the code :
device = &quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;
print(f&quot;Using {device} device&quot;)

model = GCN(1,1,1).to(device)
print(model)
# model(X).shape

loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

    
epochs = 10
for t in range(epochs):
    print(f&quot;Epoch {t+1}\n-------------------------------&quot;)
    train(train_dataloader, model, loss_fn, optimizer,Adjacency_matrix)
    test(test_dataloader, model, loss_fn,Adjacency_matrix)
print(&quot;Done!&quot;)

I got the error :

when I looking inside this ,I find the model is working well when I drop the dimension of batch-size. How I need to do to tell the model that this dimension is the batch-size which don't need to compute?


the error you're seeing is due to you trying to matrix multiple a 3d tensor (your input) by your 2D weights.
To get around this you can simply reshape your data, as we only really care about the last dim when doing matmuls:
def forward(self, input, adj):
    b_size = input.size(0)
    input = input.view(-1, input.shape[-1])
    support = torch.mm(input, self.weight)
    output = torch.spmm(adj, support)
    output = output.view(b_size,-1,output.shape[-1])
    if self.bias is not None:
        return output + self.bias
    else:
        return output

