checkpoint is different from what I posted.
batch_size_you_want = 64
max_batch_size = 8 # create dataloader with this

processed_samples = 0
for batch in dataloader:
  loss = get_loss(batch)
  loss.backward()
  processed_samples += max_batch_size
  if processed_samples >= batch_size_you_want:
    opt.step()
    opt.zero_grad()
    processed_samples = 0

This code will do training as if you were using a batch size of 64. Without ever using more memory than if you were using a batch size of 8.
checkpointing is a bit different. You will need to modify you model’s forward pass to wrap groups of operations in a torch.utils.checkpoint. This will allow you to free the memory used by the intermediary buffers that are between the ops that you gave to the checkpoint.