@staticmethod
def forward(ctx, x, shape):
    ctx.save_for_backward(x.shape)
    return x.reshape(shape)
@staticmethod
def forward(ctx, x, shape):
    ctx.save_for_backward(x.shape)
    return x.reshape(shape)