ai4co / rl4co

A PyTorch library for all things Reinforcement Learning (RL) for Combinatorial Optimization (CO)

Home Page:https://rl4.co

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

[Feature enhancement] Edge handling in GCNEnc

Junyoungpark opened this issue · comments

I found that the current implementation of GCNEnc cannot support the cases where we test the model that is not trained.
This could be critical especially when we measure the generalization performance of the trained model.

An easy fix is, inside of the forward, when the given x.shape[1] != self.edge_index[0].max() - or something more smart -, we can reconstruct the edges from the inputs.

https://github.com/kaist-silab/rl4co/blob/daf558a8ae42a4c64ba8464135c3af85423ad469/rl4co/models/nn/graph/gcn.py#LL52C39-L52C49

Good idea, added this feature in 1c22b6e.