class one_d_tensor(nn.Module):
      def __init__(self, num_classes):
        super(one_d_tensor,self).__init__()
        self.W = torch.nn.Parameter(torch.ones(num_classes,1,1))
        self.W.requires_grad = True
      def forward(self,x):
        print("self.W: ",self.W)
        mul = torch.mul(x,self.W)
        return mul

You were right! I was confused about how to deal with the unknown batch size… Simple multiplication solved it for me.