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)
def forward(self, x, y):
x = x.unsqueeze(2)
y = y.repeat(x.shape[0], x.shape[1], 1)
z = torch.cat((x, y), dim = 2)
z = self.linear(z)
z = self.softmax(z)
return z
net = Network()
C = net(A, B)
print(C.shape)
Of course you can design the model with additional layers.