Given x of shape (5, 400) and y of shape (40, 700) and a tensor d ([11,5,8,7,9]), I’m looking to repeat each row of x with each value in d so that the total length matches with y.
The below code works. But, can we do any better than using for loop in forward function?
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable

x = np.random.randn(5, 400)
x = np.array(x, np.float32)
xt = Variable(torch.from_numpy(x))

y = np.random.randn(40, 700)
y = np.array(y, np.float32)
yt = Variable(torch.from_numpy(y))

d = np.array([11,5,8,7,9], 'int32')
dt = torch.ByteTensor(d)

class model(nn.Module):
    def __init__(self, in_dim, out_dim):
        super(model, self).__init__()

        self.in_dim = in_dim
        self.out_dim = out_dim

        self.fc1 = nn.Linear(in_dim, 256)
        self.fc2 = nn.Linear(256, out_dim)


    def forward(self, x, d, n):
        h1 = self.fc1(x)
        new_h1 = h1[0].repeat(d[0], 1)
        for i in range(1, n):
            new_h1 = torch.cat((new_h1, h1[i].repeat(d[i], 1)), 0)
        h1 = new_h1
        h2 = self.fc2(h1)
        return h2

m = model(400, 700)
y_pred = m(xt, dt, len(d))
print(y_pred.size()) # 40 x 700