aa85b32ba95645f73f6fec077589c0579e4eb1b1,torch_geometric/nn/models/schnet.py,SchNet,from_qm9_pretrained,#Any#Any#Any#,85

Before Change

        idx = dataset.data.idx
        train_idx = torch.from_numpy(train_idx[np.isin(train_idx, idx)])
        val_idx = torch.from_numpy(val_idx[np.isin(val_idx, idx)])
        test_idx = torch.from_numpy(test_idx[np.isin(test_idx, idx)])

        path = osp.join(root, "trained_schnet_models", name, "best_model")
        state = torch.load(path, map_location="cpu")

After Change

        // Filter the splits to only contain characterized molecules.
        idx = dataset.data.idx
        assoc = idx.new_empty(idx.max().item() + 1)
        assoc[idx] = torch.arange(idx.size(0))

        train_idx = assoc[train_idx[np.isin(train_idx, idx)]]

Instances