torch.mm expects two matrices (2D tensors), while you seem to use two 3D tensors.
You could use torch.bmm or torch.matmul instead, which would work for these tensors.
However, usually the parameters are not depending on the batch size.
Are you sure you want to initialize them with the batch size in dim0?