Error in overlay_y_on_x

[PF-Maillard]

Hello,

I wrote a code using your project but encountered an error where considering a set of two entries X, the result of MyNet.predict(X) and MyNet.predict(X[0:1]) gives different results for the first element.

After investigation, it sees that overlay_y_on_x considers x.max() to define a label. However, if x is a unique tensor or x is a set of tensor, the result will be different (considering the max of all of the tensors). This error leads to different predictions for x if considered in a set of tensors or a unique tensor. I guess it is possible to solve it by replacing x.max to 1 or to iterate over all the tensor to get the max of each of them.