I have input tensor (x) of size (batch_size , 4 , 10). I would like to create 4 “small” separate fully connected (Linear) layers in parallel (no_inputs = 10 , no_outputs = 5).
1st linear layer will be fed by x[: , 0:1 , :]
2nd linear layer will be fed by x[: , 1:2 , :]
3rd linear layer will be fed by x[: , 2:3 , :]
4th linear layer will be fed by x[: , 3:4 , :]
I am using a single GPU for training. Is there a way to implement this in “Parallel”? So I won’t need to use a for loop as such:
output = torch.zeros(batch_size , 4 , 5)
self.layers = [nn.Linear(10 , 5) , nn.Linear(10 , 5) , nn.Linear(10 , 5) , nn.Linear(10 , 5)]
for i in range(x.shape[1]):
output[: , i : i + 1 , :] = self.layers[i](x[: , i : i+1 , :])