Best way to concatenate these tensor dimensions

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.

1 Like