🐛 Bug description

In line 165

output_list = [outputs[x * y_pred.size(0) : (x + 1) * y_pred.size(0)] for x in range(len(outputs))]

the list comprehension produces a list with len = (Batch_size * 5), where only the first 5 elements are valid and correspond to stacked [y_pred, y, y_pred * y_pred, y * y, y_pred * y] of all the batches, in cases where the batch size is greater than one the elements with index>4 are empty torch.Tensors with shape (0, C, H, W) .

Solution

This bug neither affects the output, nor consumes a lot of RAM, but I thought I should point it out.
The fix for this is pretty simple and you only need to divide the len of outputs by the batch size.

output_list = [outputs[x * y_pred.size(0) : (x + 1) * y_pred.size(0)] for x in range(int(len(outputs)/y_pred.size(0)))] # len(outputs) is B*5 so we need to divide it by B so it's only 5 -> [y_pred, y, y_pred * y_pred, y * y, y_pred * y]