@Marksein @papyan I just ran into this issue, founding this looking for confirmation, and pretty sure I fixed the bug
change the original code here…
    def call_checkpoint_bottleneck(self, input):
        # type: (List[Tensor]) -> Tensor
        def closure(*inputs):
            return self.bn_function(*inputs)

        return cp.checkpoint(closure, input)

to
    def call_checkpoint_bottleneck(self, input):
        # type: (List[Tensor]) -> Tensor
        def closure(*inputs):
            return self.bn_function(inputs) # remove *

        return cp.checkpoint(closure, *input) # add *

Very subtle!