Fixing AffineAdapterSigmoid for 2 dimensional inputs

[silvandeleemput]

Thank you for your reply, I modified it according to your description, but it brought another error.

If I use AffineAdapterNaive instead of AffineAdapterSigmoid, it works.

Btw, if you still want to use the AffineAdapterSigmoid with the AffineCoupling, you can use the following modified AffineAdapterSigmoid implemention, which should fix your error:

class AffineAdapterSigmoidModified(torch.nn.Module):
    """ Sigmoid based affine adapter, modified to work for 2 dimensional outputs

        Partitions the output h of f(x) = h into s and t by extracting every odd and even channel
        Outputs sigmoid(s), t
    """
    def __init__(self, module):
        super(AffineAdapterSigmoidModified, self).__init__()
        self.f = module

    def forward(self, x):
        h = self.f(x)
        assert h.shape[1] % 2 == 0  # nosec
        scale = torch.sigmoid(h[:, 1::2] + 2.0)
        shift = h[:, 0::2]
        return scale, shift

Originally posted by @silvandeleemput in #63 (comment)