Update PolyFusion/SchNet.py

PolyFusion/SchNet.py

@@ -1,10 +1,6 @@
 """
 SchNet.py
 SchNet-based masked pretraining on polymer conformer geometry.
-
-This file provides (and uses internally):
-- NodeSchNetWrapper (used by CL.py AND used internally by MaskedSchNet)
-- MaskedSchNet training code (behavior preserved)
 """
 
 from __future__ import annotations
@@ -276,12 +272,12 @@ class NodeSchNet(nn.Module):
 
 
 # =============================================================================
-# Wrapper used by CL.py AND used internally by MaskedSchNet
+# Wrapper used by MaskedSchNet
 # =============================================================================
 
 class NodeSchNetWrapper(nn.Module):
     """
-    - Produces pooled embedding for CL (mean pooling + pool_proj)
+    - Produces pooled embedding (mean pooling + pool_proj)
     - Provides node_logits(...) for reconstruction
     """
 
@@ -619,7 +615,7 @@ def train_and_eval(args: argparse.Namespace) -> None:
         except Exception as e:
             print(f"\nFailed to load best model from {best_model_path}: {e}")
 
-    # Final evaluation (same as your original)
+    # Final evaluation
     model.eval()
     preds_z_all, true_z_all = [], []
     pred_dists_all, true_dists_all = [], []