def METHOD_NAME(self):
    self.assertEqual(find_labels(TFBertForSequenceClassification), ["labels"])
    self.assertEqual(find_labels(TFBertForPreTraining), ["labels", "next_sentence_label"])
    self.assertEqual(find_labels(TFBertForQuestionAnswering), ["start_positions", "end_positions"])
    # find_labels works regardless of the class name (it detects the framework through inheritance)
    class DummyModel(TFBertForSequenceClassification):
        pass
    self.assertEqual(find_labels(DummyModel), ["labels"])