yolact_edge: Custom YOLACT model trains well but shows IndexError in case of eval.py

Hi,

Thank you for developing this amazing tool.

I recently modified the yolact.py and added one more mask so that the network predicts 2 of them (instead of 1). I am also using a custom dataset. The training works fine but eval.py is showing the following error:

File "eval.py", line 173, in prep_display
    masks = t[3][idx]
IndexError: too many indices for tensor of dimension 3

I checked the size of idx and t[3] to debug the issue and found the following information:

idx: tensor([0, 1, 2, 3, 4])
idx.size(): torch.Size([5])
t[3].size(): torch.Size([100, 480, 640])

It seems that the error is misleading. Anyway, after a while, I realized that in your eval.py, the variable idx is simply :args.top_k. So I changed it but then the output from eval.py was shown incorrectly.

Moreover, the statement masks = t[3][idx] works fine if PredictionModuleTRTWrapper is disabled (just comment 2 lines). I should also mention that the network shows the above error in forward propagation step in eval.py. This leads me to think that the pred_layer is not set properly.

Can you please help me out? Is input_sizes variable is set properly or needs some modification?

Thank you so much.