modify a bit

inference.py

@@ -38,7 +38,7 @@ device = torch.device(f"cuda:{gpu}" if gpu>=0 else "cpu")
 
 #BERTモデルの定義
 class BertClassifier(nn.Module, PyTorchModelHubMixin):
-    def __init__(self, cls_num: int):
+    def __init__(self, cls_num: int = 1):
         super().__init__()
         self.bert = T.BertModel.from_pretrained(bert_model_name, output_attentions=True)
         self.fc = nn.Linear(768, cls_num, bias=True)
@@ -124,6 +124,7 @@ def make_traind_model():
         # trained_models.append(trained_model)
         model_name = MODEL_BASE + str(k)
         trained_model = BertClassifier.from_pretrained(model_name).to(device)
+        print(f"Got model {model_name}")
         trained_models.append(trained_model)
     return trained_models