topobench.nn.wrappers.combinatorial package

Submodules

topobench.nn.wrappers.combinatorial.tune_wrapper module

Wrapper for the TopoTune model.

class topobench.nn.wrappers.combinatorial.tune_wrapper.TuneWrapper(backbone, **kwargs)

Bases: AbstractWrapper

Wrapper for the TopoTune model.

This wrapper defines the forward pass of the TopoTune model.

forward(batch)

Forward pass for the Tune wrapper.

Parameters:

batchtorch_geometric.data.Data

Batch object containing the batched data.

Returns:

dict

Dictionary containing the updated model output.

Module contents

Wrappers for combinatorial neural networks with automated exports.

class topobench.nn.wrappers.combinatorial.TuneWrapper(backbone, **kwargs)

Bases: AbstractWrapper

Wrapper for the TopoTune model.

This wrapper defines the forward pass of the TopoTune model.

forward(batch)

Forward pass for the Tune wrapper.

Parameters:

batchtorch_geometric.data.Data

Batch object containing the batched data.

Returns:

dict

Dictionary containing the updated model output.