I am not absolutely sure in this point, but do you really want to have the batch_size as one Dimension of your weight w?
Shouldn’t the weight be the same for every element in the batch and you just expand it in the forward computation like this:
class MnistModel(nn.Module):
    def __init__(self, batch_size):
        super(MnistModel, self).__init__()
        self.w = torch.nn.Parameter(torch.empty(784, 10).uniform_(0, 1))
        self.b = torch.nn.Parameter(torch.empty(10).uniform_(0, 1))

    def forward(self, x):
        exp_w = self.w.expand(x.size(0), self.w.size(0), self.w.size(1))
        return torch.bmm(x, exp_w) + self.b

?