I have a Siamese Network with a triplet loss function at the end. E.g.:
> class Siamese(nn.Module):
>     def __init__(self, ae_net):
>         super(Siamese, self).__init__()
>         self.ae_net = ae_net
> 
>     def forward(self, x1, x2, x3, hidden):
>         
>         a = self.ae_net(x1, hidden)
>         b = self.ae_net(x2, hidden)
>         c = self.ae_net(x3, hidden)
>         return a, b, c

The network that repeats itself (i.e., self.ae_net) is an LSTM with inputs of varying lengths, so I’m not sure I can use nn.DataParallel.  I was wondering if there was a way to assign the implementation of every instance of self.ae_net to a different GPU, so that they would be calculated in parallel.
Thanks