Update PolyFusion/CL.py

PolyFusion/CL.py

@@ -28,10 +28,10 @@ import torch.nn.functional as F
 from torch.utils.data import Dataset, DataLoader
 
 # Shared model utilities 
-from GINE import GineBlock, MaskedGINE, match_edge_attr_to_index, safe_get
-from SchNet import NodeSchNetWrapper
-from Transformer import PooledFingerprintEncoder as FingerprintEncoder
-from DeBERTav2 import PSMILESDebertaEncoder, build_psmiles_tokenizer
+from PolyFusion.GINE import GineBlock, MaskedGINE, match_edge_attr_to_index, safe_get
+from PolyFusion.SchNet import NodeSchNetWrapper
+from PolyFusion.Transformer import PooledFingerprintEncoder as FingerprintEncoder
+from PolyFusion.DeBERTav2 import PSMILESDebertaEncoder, build_psmiles_tokenizer
 
 # HF Trainer & Transformers
 from transformers import TrainingArguments, Trainer