For version 2.0 i.e, original paper 2D.ipynb gives error

[aknirala]

I am trying to understand the original code as understanding original paper seems simpler. While trying to optimize Robust Neural Network in 2D.ipynb I get error as:

IndexError                                Traceback (most recent call last)
Input In [7], in <cell line: 19>()
     18 opt = optim.Adam(robust_net.parameters(), lr=1e-3)
     19 for i in range(1000):
---> 20     robust_ce, robust_err = robust_loss(robust_net, epsilon, X, y)
     21     out = robust_net(X)
     22     l2 = nn.CrossEntropyLoss()(out, Variable(y))

File ~/PhD/Compo/from_git/convex_adversarial/examples/../convex_adversarial/dual.py:189, in robust_loss(net, epsilon, X, y, size_average, alpha_grad, scatter_grad)
    187 if X.is_cuda:
    188     c = c.cuda()
--> 189 f = -dual.g(c)
    190 err = (f.data.max(1)[1] != y.data)
    191 if size_average: 

File ~/PhD/Compo/from_git/convex_adversarial/examples/../convex_adversarial/dual.py:162, in DualNetBounds.g(self, c)
    160 print("self.I_neg[i-1].shape: ", self.I_neg[i-1].shape)
    161 print("self.I_neg[i-1].unsqueeze(1).shape: ", self.I_neg[i-1].unsqueeze(1).shape)
--> 162 out[self.I_neg[i-1].unsqueeze(1)] = 0
    163 if not self.I_empty[i-1]:
    164     if self.alpha_grad: 

IndexError: The shape of the mask [12, 1, 100] at index 1 does not match the shape of the indexed tensor [12, 2, 100] at index 1

I added few print statements in dual.py to debug.
I also changed the import statement from:
from convex_adversarial import robust_loss_batch
to
from convex_adversarial import robust_loss_batch
as it was not able to find robust_loss_batch

Thanks in advance for any help.