fixed evaluate function

2024-10-21 20:12:20 +00:00 · 2024-10-21 20:12:20 +00:00 · 3f116b583d
commit 3f116b583d
parent 240c16b495
1 changed files with 1 additions and 0 deletions
--- a/code/src/train_codebert_mlm.py
+++ b/code/src/train_codebert_mlm.py
@ -208,6 +208,7 @@ def evaluate(model: PreTrainedModel, dataloader: DataLoader) -> Tuple[float, flo
    
    avg_loss: float = total_loss / len(dataloader)
    avg_acc: float = total_acc / len(dataloader.dataset)
+    model.train()
    return avg_loss, avg_acc

 def main() -> None: