search

prateekjoshi565 / Fine-Tuning-BERT

Apache License 2.0
139 stars 119 forks source link

Extend for multi class classification #1

Open MainakMaitra opened 3 years ago

MainakMaitra commented 3 years ago

Hi

Can you please help how to extend it for multi-class classification?

Thanks in advance!

prateekjoshi565 commented 3 years ago

Hi, simply replace the output units from 2 to the number of classes in your dataset in the layer below:

self.fc2 = nn.Linear(512,2)

MainakMaitra commented 3 years ago

Hi, thank you for responding so quickly, and thanks for the help!