I’m training a network in which the first layer is a filter bank like projection and the second is a convolution. The example below is just to illustrate the problem.
When using this network inside of DataParallel with more than one GPU I run into “tensors are on different GPUs”.
I also tried to not explicitly move the tensors to the GPU with the hope that calling .cuda() on model would solve the issue.
@SimonW Any thoughts on this?
Does anyone have a solution to this problem?
model = DataParallel(Network(frame_size, stride).cuda())
class Network(torch.nn.Module):
    def __init__(self, frame_size, stride):
        super(Network, self).__init__()
        self.projection = Projection(frame_size, stride)
        self.conf = torch.nn.Conv1d(frame_size/2, 1, 2, 1)

    def forward(self, signal):
        output = self.projection(signal)
        output = self.conv(output)
        return output

class Projection(torch.autograd.Function):
    def __init__(self, frame_size, stride):
        super(Projection, self).__init__()
        x = np.array(range(0, frame_size))
        f = np.array([np.sin(x*2*f*np.pi/frame_size) for f in range(0,int(frame_size/2))])
        f = f.astype(float)
        f = torch.from_numpy(f)
        f = torch.unsqueeze(f, 1).float()
        self.conv = torch.nn.Conv1d(1, int(frame_size/2),
                                    kernel_size=frame_size,
                                    stride = stride,
                                    bias=False)
        self.conv.weight.data = f
        if torch.cuda.is_available():
            self.conv.weight.data = self.conv.weight.data.cuda()
        self.conv.weight.requires_grad = False

    def forward(self, signals):
        signals = Variable(signals)
        conv_signal = self.conv(signals).data
        return conv_signal


class Filterbank(torch.nn.Module):
    def __init__(self, frame_size, stride):
        super(Filterbank, self).__init__()
        self.proj = Projection(frame_size, stride)

    def forward(self, signals):
        signals = torch.unsqueeze(signals, 1)
        signals_proj = self.proj(signals)**2
        return sin_signals