Captum IG problem when applied to Graph WaveNet; Issue with mismatch in input and output dim becuase of Latent Embeddings During Training

[gcalistanford]

I am applying Captum Integrated Gradients tool to understand my GraphWaveNet model based on https://github.com/nnzhan/Graph-WaveNet/tree/master

The problem is that I continue to get an error RuntimeError: Given groups=1, weight of size [64, 10, 1, 1], expected input[1, 12, 1596, 37] to have 10 channels, but got 12 channels instead. This to my understanding is because during training of the model I have some internal adjacency matrix learning not part of the original model inputs.

Here is my model architecture.
gwnet(
(filter_convs): ModuleList(
(0): Conv2d(64, 64, kernel_size=(1, 2), stride=(1, 1))
(1): Conv2d(64, 64, kernel_size=(1, 2), stride=(1, 1), dilation=(2, 2))
(2): Conv2d(64, 64, kernel_size=(1, 2), stride=(1, 1))
(3): Conv2d(64, 64, kernel_size=(1, 2), stride=(1, 1), dilation=(2, 2))
(4): Conv2d(64, 64, kernel_size=(1, 2), stride=(1, 1))
(5): Conv2d(64, 64, kernel_size=(1, 2), stride=(1, 1), dilation=(2, 2))
(6): Conv2d(64, 64, kernel_size=(1, 2), stride=(1, 1))
(7): Conv2d(64, 64, kernel_size=(1, 2), stride=(1, 1), dilation=(2, 2))
)
(gate_convs): ModuleList(
(0): Conv1d(64, 64, kernel_size=(1, 2), stride=(1,))
(1): Conv1d(64, 64, kernel_size=(1, 2), stride=(1,), dilation=(2,))
(2): Conv1d(64, 64, kernel_size=(1, 2), stride=(1,))
(3): Conv1d(64, 64, kernel_size=(1, 2), stride=(1,), dilation=(2,))
(4): Conv1d(64, 64, kernel_size=(1, 2), stride=(1,))
(5): Conv1d(64, 64, kernel_size=(1, 2), stride=(1,), dilation=(2,))
(6): Conv1d(64, 64, kernel_size=(1, 2), stride=(1,))
(7): Conv1d(64, 64, kernel_size=(1, 2), stride=(1,), dilation=(2,))
)
(residual_convs): ModuleList(
(0): Conv1d(64, 64, kernel_size=(1, 1), stride=(1,))
(1): Conv1d(64, 64, kernel_size=(1, 1), stride=(1,))
(2): Conv1d(64, 64, kernel_size=(1, 1), stride=(1,))
(3): Conv1d(64, 64, kernel_size=(1, 1), stride=(1,))
(4): Conv1d(64, 64, kernel_size=(1, 1), stride=(1,))
(5): Conv1d(64, 64, kernel_size=(1, 1), stride=(1,))
(6): Conv1d(64, 64, kernel_size=(1, 1), stride=(1,))
(7): Conv1d(64, 64, kernel_size=(1, 1), stride=(1,))
)
(skip_convs): ModuleList(
(0): Conv1d(64, 512, kernel_size=(1, 1), stride=(1,))
(1): Conv1d(64, 512, kernel_size=(1, 1), stride=(1,))
(2): Conv1d(64, 512, kernel_size=(1, 1), stride=(1,))
(3): Conv1d(64, 512, kernel_size=(1, 1), stride=(1,))
(4): Conv1d(64, 512, kernel_size=(1, 1), stride=(1,))
(5): Conv1d(64, 512, kernel_size=(1, 1), stride=(1,))
(6): Conv1d(64, 512, kernel_size=(1, 1), stride=(1,))
(7): Conv1d(64, 512, kernel_size=(1, 1), stride=(1,))
)
(bn): ModuleList(
(0): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(5): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(6): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(7): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(gconv): ModuleList(
(0): gcn(
(nconv): nconv()
(mlp): linear(
(mlp): Conv2d(320, 64, kernel_size=(1, 1), stride=(1, 1))
)
)
(1): gcn(
(nconv): nconv()
(mlp): linear(
(mlp): Conv2d(320, 64, kernel_size=(1, 1), stride=(1, 1))
)
)
(2): gcn(
(nconv): nconv()
(mlp): linear(
(mlp): Conv2d(320, 64, kernel_size=(1, 1), stride=(1, 1))
)
)
(3): gcn(
(nconv): nconv()
(mlp): linear(
(mlp): Conv2d(320, 64, kernel_size=(1, 1), stride=(1, 1))
)
)
(4): gcn(
(nconv): nconv()
(mlp): linear(
(mlp): Conv2d(320, 64, kernel_size=(1, 1), stride=(1, 1))
)
)
(5): gcn(
(nconv): nconv()
(mlp): linear(
(mlp): Conv2d(320, 64, kernel_size=(1, 1), stride=(1, 1))
)
)
(6): gcn(
(nconv): nconv()
(mlp): linear(
(mlp): Conv2d(320, 64, kernel_size=(1, 1), stride=(1, 1))
)
)
(7): gcn(
(nconv): nconv()
(mlp): linear(
(mlp): Conv2d(320, 64, kernel_size=(1, 1), stride=(1, 1))
)
)
)
(start_conv): Conv2d(10, 64, kernel_size=(1, 1), stride=(1, 1))
(end_conv_1): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1))
(end_conv_2): Conv2d(1024, 12, kernel_size=(1, 1), stride=(1, 1))
)

Does anyone know or can direct me to a solution for this? It is quite frustrating and I assume a solution exists somewhere I just cannot find it right now.

Best,

Gabe