Class Acts::TorchEdgeClassifier

class TorchEdgeClassifier : public Acts::EdgeClassificationBase

Public Functions

TorchEdgeClassifier(const Config &cfg)
~TorchEdgeClassifier()
inline Config config() const
virtual std::tuple<std::any, std::any, std::any> operator()(std::any nodes, std::any edges, const Logger &logger) override

Perform edge classification.

Parameters

  • nodes – Node tensor with shape (n_nodes, n_node_features)

  • edges – Edge-index tensor with shape (2, n_edges)

  • logger – Logger instance

Returns

(node_tensor, edge_tensor, score_tensor)

struct Config

Public Members

float cut = 0.21
std::string modelPath
int nChunks = 1