Runtime Error While Training PixelSNAIL

[barisbatuhan]

trafficstars

Hi,

I am training the PixelSNAIL network with 128x128 image sizes and currently having the runtime error below:

background = self.background[:, :, :height, :].expand(batch, 2, height, width)
RuntimeError: The expanded size of the tensor (16) must match the existing size (32) at
non-singleton dimension 3.  Target sizes: [64, 2, 16, 16].  Tensor sizes: [1, 2, 16, 32]

The problem is caused by the the line 408 in pixelsnail.py:

background = self.background[:, :, :height, :].expand(batch, 2, height, width)

If I change the line to:

background = self.background[:, :, :height, :width].expand(batch, 2, height, width)

the code works without any errors but I am not sure if that is the correct way to fix the error. How can this error be fixed?

[rosinality]

I think you can change shape argument of PixelSNAIL.