From ce8f7bb1d687cd6762e071cea21c5e281ac1ce46 Mon Sep 17 00:00:00 2001 From: Evening Date: Wed, 5 Jun 2024 17:26:14 +0800 Subject: [PATCH] Reduce scope of MixMatch custom criteria saving --- src/frdc/train/mixmatch_module.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/frdc/train/mixmatch_module.py b/src/frdc/train/mixmatch_module.py index fb24b9b..9617751 100644 --- a/src/frdc/train/mixmatch_module.py +++ b/src/frdc/train/mixmatch_module.py @@ -266,7 +266,7 @@ def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None: save_unfrozen( self, checkpoint, - include_also=lambda k: k.startswith("_ema_model."), + include_also=lambda k: k.startswith("_ema_model.fc."), ) def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None: