A quick question about torch.utils.checkpoint.
My classification head is:
class ClassificationHead(nn.Module):
  def __init__(self, config):
    super().__init__()
    self.dense = nn.Linear(config.hidden_size, config.hidden_size)
    self.dropout = nn.Dropout(config.hidden_dropout_prob)
    self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
 forward(self, inputs):
   ....

I use it in another bigger model:
sf_logits = checkpoint(self.sf_classifier, sequence_output)

But I’m getting an error: “TypeError: ‘module’ object is not callable”.
In the documentation, it doesn’t say what is the type of the function to be passed.