Solved by ptrblck in post #2 
              
              
                You can add additional dimensions like this: 
batch_size = 10
add_features = 5
in_features = 7
out_features = 12

x = torch.randn(batch_size, add_features, in_features)
lin = nn.Linear(in_features, out_features)
output = lin(x)

Is this what you are looking for?