How does torch differentiate between batch and single datasets?

Here’s my neural network.

class NeuralNetwork(nn.Module):
    def __init__(self, state_size, action_size):
        super(NeuralNetwork, self).__init__()
        self.state_size = state_size
        self.action_size = action_size
#         self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(1, 30),
            nn.ReLU(),
            nn.Linear(30, 30),
            nn.ReLU(),
            nn.Linear(30, action_size)
        )
    def forward(self, x):
        x = self.linear_relu_stack(x)
        return x

As you can imagine, it can compute tensor inputs of the shape torch.Size([1]). However, when I try to feed it batch data, for instance, shape torch.Size([10]) it throws the following error -

RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x10 and 1x30)

For instance, this code works -
net = NeuralNetwork(10, 5)
x1 = torch.rand(1)
print(x1.shape)
out = net(x1)

But this fails -

x2 = torch.rand(10)
print(x2.shape)
out = net(x2)

In your nn.Sequential object you have a single input node which is why x1 = torch.rand(1) works. However, when you pass x1 = torch.rand(10) you’re effectively trying to do matmul between a vector of length 1 and vector of length 10 which you can’t do. So, you need to include the batch size and the input size.

So change x1 = torch.rand(10, 1) and it’ll calculate the output for a batch of size 10.

2 Likes