Fix a bug in unfreeze (#660)

*Description of changes:*
When we unfreeze the LM model, we don't need to clean up the cache if
the number of train nodes is 0.

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

Co-authored-by: Ubuntu <[email protected]>

python/graphstorm/model/lm_embed.py

@@ -516,10 +516,11 @@ def freeze(self, _):
     def unfreeze(self):
         """ Disable Bert caching
         """
-        self.use_cache = False
-        # We should clear up the LM cache here. When someone calls unfreeze,
-        # we expect that the LM model will be fine-tuned.
-        self.lm_emb_cache.clear_cache()
+        if self.num_train > 0:
+            self.use_cache = False
+            # We should clear up the LM cache here. When someone calls unfreeze,
+            # we expect that the LM model will be fine-tuned.
+            self.lm_emb_cache.clear_cache()
 
     def require_cache_embed(self):
         """ Whether to cache the embeddings for inference.
@@ -726,10 +727,11 @@ def freeze(self, _):
     def unfreeze(self):
         """ Disable Bert caching
         """
-        self.use_cache = False
-        # We should clear up the LM cache here. When someone calls unfreeze,
-        # we expect that the LM model will be fine-tuned.
-        self.lm_emb_cache.clear_cache()
+        if self.num_train > 0:
+            self.use_cache = False
+            # We should clear up the LM cache here. When someone calls unfreeze,
+            # we expect that the LM model will be fine-tuned.
+            self.lm_emb_cache.clear_cache()
 
     def require_cache_embed(self):
         """ Whether to cache the embeddings for inference.