Here is a small example:
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.hidden = nn.Linear(1, 10)
        self.predict = nn.Linear(10, 1)
        
    def forward(self, x):
        x = F.relu(self.hidden(x))
        x = self.predict(x)
        return x


model = MyModel()

batch_size = 10
x = torch.randn(batch_size, 1)

output = model(x)
print(output.shape)
> torch.Size([10, 1])

output = model(x[0]) # Only first sample
print(output.shape)
> torch.Size([1])

Is this what you are looking for?