I have it like this
import torch
bla bla bla
mels_self = torch.utils.checkpoint.checkpoint(self.decoder_self_run, mels, encoder_outputs, batch['text_lengths'], batch['mel_lengths'])

With
from torch.utils.checkpoint import checkpoint

same code works fine, why so?