Skip to content

Commit

Permalink
Fixed molecule data loader to be able to take , , and a composition o…
Browse files Browse the repository at this point in the history
…bject with transforms.
  • Loading branch information
CoopLo committed Mar 25, 2024
1 parent b24752c commit d122d36
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion auglichem/molecule/data/_molecule_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,8 @@ def __init__(self, dataset, data_path=None, transform=None, smiles_data=None, la
transform = Compose(transform)
elif(not(isinstance(transform, Compose)) and (transform is not None)):
transform = Compose([transform])
elif(transform is None):
transform = Compose([])
self.transform = transform

if(smiles_data is None):
Expand Down Expand Up @@ -134,7 +136,7 @@ def __init__(self, dataset, data_path=None, transform=None, smiles_data=None, la
def _handle_motifs(self):
# MotifRemoval adds multiple new SMILES strings to our data, and must be done
# upon training set initialization
if(not self.test_mode):
if(self._training_set):
if(self._train_warn): # Catches if not set through get_data_loaders()
raise ValueError(
"_training_set is for internal use only. " + \
Expand Down

0 comments on commit d122d36

Please sign in to comment.