I have a BERT model for Classification which I am using for classification task. The Model Structure is as follows:
class BERT_Subject_Classifier(nn.Module):