It actually didn’t work. It sure can input many input as the *args, but list output causes error. For example, my output was like,
return [tensor1, tensor2, tensor3, tensor4, tensor4]

and the error would be like,
TypeError: CheckpointFunctionBackward.forward: expected Variable (got list) for return value 0

I hope I am using it right.