ardigen / MAT

The official implementation of the Molecule Attention Transformer.
MIT License
234 stars 57 forks source link

mol_collate_func potential bug #10

Closed genec1 closed 4 years ago

genec1 commented 4 years ago

I've modified the data_loader code and am running into a crash in featurization.data_utils.mol_collate_func here:

   for molecule in batch:
         if type(molecule.y[0]) == np.ndarray:
            labels.append(molecule.y[0])
        else:
            labels.append(molecule.y)

I didn't debug the original code to see why it doesn't crash here, but with my code molecule.y is a scalar, so molecule.y[0] doesn't work. It seems like it should be:

   for molecule in batch:
         if type(molecule.y) == np.ndarray:
            labels.append(molecule.y[0])
        else:
            labels.append(molecule.y)
Mazzza commented 4 years ago

Hi,

In our original code, molecule.y is always a list, so that one can always call molecule.y[0].

This if statement that you show here differentiates between multi-task (where molecule.y[0] is np.ndarray) and single-task setting (where molecule.y[0] is a scalar).