/torchdrug

name: torchdrug
description: Build and troubleshoot TorchDrug 0.2.1 workflows for molecular graphs, property prediction, self-supervised pretraining, molecule generation, retrosynthesis, protein representation learning, and knowledge graph reasoning. Use when code imports torchdrug or needs its datasets, models, tasks, or Engine.
license: Apache-2.0 license
compatibility: TorchDrug 0.2.1 requires Python 3.7-3.10 and supports PyTorch 1.8-2.0. Apple Silicon is CPU-only; MPS is unsupported.
allowed-tools: Read Write Edit Bash
metadata:
  version: "1.2"
  skill-author: K-Dense Inc.

TorchDrug

Use TorchDrug as a modular PyTorch graph-learning stack:

1. load a `datasets.*` dataset, 2. choose a `models.*` representation model, 3. wrap it in a `tasks.*` objective, 4. train and evaluate it with `core.Engine`.

The current official documentation and latest release are both **0.2.1**. Treat newer Python or PyTorch combinations as unverified rather than silently assuming compatibility.

Start with the version guard

Before generating or debugging code, inspect the environment:

python --version
python -c "import torch; print(torch.__version__)"
python -c "import torchdrug; print(torchdrug.__version__)"

The supported matrix for TorchDrug 0.2.1 is:

• Python 3.7 through 3.10
• PyTorch 1.8 through 2.0
• Linux, Windows, or macOS
• Apple Silicon: PyTorch 1.13 or later, CPU only; no MPS support

If the project uses Python 3.11+ or PyTorch 2.1+, create a compatible environment or explicitly test a source build. Do not present such combinations as supported.

Installation

Prefer a dedicated Python 3.10 environment and pin the TorchDrug release:

uv venv --python 3.10
source .venv/bin/activate
uv pip install "torch==2.0.0"

Install `torch-scatter` and `torch-cluster` wheels matched to the exact PyTorch and CUDA pair, following the [official installation page](https://torchdrug.ai/docs/installation.html). For a CPU-only PyTorch 2.0 environment, one reproducible wheel combination is:

uv pip install "torch-scatter==2.1.1" "torch-cluster==1.6.1" \
  --find-links "https://data.pyg.org/whl/torch-2.0.0+cpu.html"
uv pip install "torchdrug==0.2.1"

Do not copy a CUDA wheel URL between environments. Match the PyTorch version, CUDA build, Python ABI, and platform. On Apple Silicon, the official docs require building `torch-scatter` and `torch-cluster` from source; pin reviewed source revisions and expect CPU execution.

Canonical property-prediction workflow

Use the documented ClinTox → GIN → `PropertyPrediction` → `Engine` pattern:

import torch
from torchdrug import core, datasets, models, tasks

dataset = datasets.ClinTox("~/molecule-datasets/")
lengths = [int(0.8 * len(dataset)), int(0.1 * len(dataset))]
lengths.append(len(dataset) - sum(lengths))
train_set, valid_set, test_set = torch.utils.data.random_split(dataset, lengths)

model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[256, 256, 256, 256],
    short_cut=True,
    batch_norm=True,
    concat_hidden=True,
)
task = tasks.PropertyPrediction(
    model,
    task=dataset.tasks,
    criterion="bce",
    metric=("auprc", "auroc"),
)

optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver = core.Engine(
    task,
    train_set,
    valid_set,
    test_set,
    optimizer,
    batch_size=1024,
)
solver.train(num_epoch=100)
solver.evaluate("valid")

Add `gpus=[0]` only when a supported CUDA device is available. Omit `gpus` for CPU execution.

For binary classification, `task.predict(batch)` returns logits; apply `torch.sigmoid` when probabilities are needed. In 0.2.1, normalized regression predictions are returned on the original target scale, which is a breaking change from older releases.

Choose the official workflow

Molecular property prediction

• Dataset: `datasets.ClinTox`, `BBBP`, `Tox21`, `QM9`, or another documented

molecule dataset.

• Model: start with `models.GIN`; use `edge_input_dim` when the selected feature

configuration supplies edge features.

• Task: `tasks.PropertyPrediction`.
• Read [molecular property prediction](references/molecular_property_prediction.md).

Self-supervised molecular pretraining

• InfoGraph: `models.InfoGraph(gin_model, separate_model=False)` wrapped by

`tasks.Unsupervised`.

• Attribute masking: `tasks.AttributeMasking(model, mask_rate=0.15)`.
• Recreate the same encoder for fine-tuning, then load the checkpoint with

`strict=False` before training `tasks.PropertyPrediction`.

• Read [molecular property prediction](references/molecular_property_prediction.md).

Molecule generation

• Dataset: `datasets.ZINC250k(..., kekulize=True, atom_feature="symbol")`.
• GCPN: an `models.RGCN` encoder wrapped by `tasks.GCPNGeneration`.
• GraphAF: node and edge `models.GraphAF` flows wrapped by

`tasks.AutoregressiveGeneration`.

• Supported optimization tasks in the tutorial are `"qed"` and `"plogp"`;

criteria are `"nll"` and/or `"ppo"`.

• Read [molecular generation](references/molecular_generation.md).

Retrosynthesis

• Create two synchronized `datasets.USPTO50k` views: reaction mode for center

identification and `as_synthon=True` for synthon completion.

• Train `tasks.CenterIdentification` and `tasks.SynthonCompletion` separately.
• Combine the trained tasks with `tasks.Retrosynthesis`; do not pass raw models

directly to the end-to-end task.

• Read [retrosynthesis](references/retrosynthesis.md).

Knowledge graph reasoning

• Embedding workflow: `datasets.FB15k237` → `models.RotatE` →

`tasks.KnowledgeGraphCompletion`.

• Neural reasoning workflow: `models.NeuralLP` with `fact_ratio=0.75`.
• Read [knowledge graph reasoning](references/knowledge_graphs.md).

Protein modeling

• Build proteins with `data.Protein.from_sequence`, `from_pdb`, or

`from_molecule`.

• Sequence encoders include `models.ESM`, `ProteinCNN`, `ProteinResNet`,

`ProteinLSTM`, and `ProteinBERT`; structure encoders include