But if I take the original code and remove the call to checkpoint() and just do x = self.run_func(self.conv1)(x, dummy_input). You still get the exact same result right? So the issue is not with checkpoint I think. Can you confirm?