Yes I use save for backward
    @staticmethod
     def forward(ctx, mu, sigma, landa, eps, dimX, dimY, dimZ):

    # import pdb
    # pdb.set_trace()

    ctx.save_for_backward(mu, sigma, landa, eps, dimX, dimY, dimZ)

dimX, dimY and dimZ are integers. But as you can see mu, sigma, landa, and eps are input of my Function. So I should be allowed to save them for backward