Solved by ksz in post #6 
              
              
                You can use this: 

A = torch.rand(30, 20)
B = torch.rand(10)

class Network(torch.nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        
        self.linear = torch.nn.Linear(11, 10) #11 -> 1 from x + 10 from y
        self.softmax = torch.nn.Softmax(dim=1)
        
  …