StyleGAN2 network error

[danny0628]

When training with custom .png image, the following error occurs.

class Smooth(nn.Module):
    def __init__(self):
        super().__init__()
        kernel = [[1, 2, 1], [2, 4, 2], [1, 2, 1]]  # define the blue kernel
        kernel = torch.tensor([[kernel]], dtype=torch.float32)  # making it a pytorch tensor

        kernel /= kernel.sum()  # Normalize the kernel
        self.kernel = nn.Parameter(kernel, requires_grad=False)  # fix the kernel so it doesn't get update
        self.pad = nn.ReplicationPad2d(1)  # pad the kernel

    def forward(self, x):
        print(x.shape)
        b, c, h, w = x.shape
        x = x.view(-1, 1, h, w)  # reshape
        x = self.pad(x)  # pad
        x = F.conv2d(x, self.kernel)  # smooth
        return x.view(b, c, h, w)  # reshape back and return

File "/home/mjkim1/workspace/fl_learning/flower/model.py", line 314, in forward b, c, h, w = x.shape ValueError: not enough values to unpack (expected 4, got 3)

So I try to debuging, and I could see the x.shape changed. How can I solve this problem?