def pixel_shuffle(self, x):
    out = self.ps(x)
    b = out.shape[0]
    c = out.shape[1]
    h = out.shape[2]
    w = out.shape[3]
    out = F.reshape(out, (b, 2, 2, (c // 4), h, w))
    out = F.transpose(out, (0, 3, 4, 1, 5, 2))
    out = F.reshape(out, (b, (c // 4), (h * 2), (w * 2)))
    return out
def pixel_shuffle(self, x):
    out = self.ps(x)
    b = out.shape[0]
    c = out.shape[1]
    h = out.shape[2]
    w = out.shape[3]
    out = F.reshape(out, (b, 2, 2, (c // 4), h, w))
    out = F.transpose(out, (0, 3, 4, 1, 5, 2))
    out = F.reshape(out, (b, (c // 4), (h * 2), (w * 2)))
    return out