Wasserstein GAN dot product

1

I am trying to modify the Wasserstain GAN code to incorporate data from more than one source. The original code is here:

https://github.com/martinarjovsky/WassersteinGAN/blob/master/main.py

I run into problems computing the dot product of two vectors.

# train with fake
noise.resize_(opt.batchSize, nz, 1, 1).normal_(0, 1)
noisev = Variable(noise)
fake = Variable(netG_image(torch.bmm(netG_audio(inputv_audio).data, noisev)).data)

I get the following error:

#RuntimeError: Expected 3-dimensional tensor, but got 1-dimensional tensor for argument #1 'batch1' (while checking arguments for bmm)

It appears that netG_audio (an instance of MLP_D) returns one vector instead of multiple vectors.

class MLP_D(nn.Module):
    def __init__(self, isizex, isizey, nz, nc, ndf, ngpu):
        super(MLP_D, self).__init__()
        self.ngpu = ngpu

        main = nn.Sequential(
            # Z goes into a linear of size: ndf
            nn.Linear(nc * isizex * isizey, ndf),
            nn.ReLU(True),
            nn.Linear(ndf, ndf),
            nn.ReLU(True),
            nn.Linear(ndf, ndf),
            nn.ReLU(True),
            nn.Linear(ndf, ndf),
        )
        self.main = main
        self.nc = nc
        self.isizex = isizex
        self.isizey = isizey
        self.nz = nz

    def forward(self, input):
        input = input.view(input.size(0),
                           input.size(1) * input.size(2) * input.size(3))
        if torch.cuda.is_available() and isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
        else:
            output = self.main(input)
        output = output.mean(0)
        return output

However, in the original code netG(noisev) seems to return a batch of vectors. What am I missing?