bugs in LossesHelper.removeSqueezableDimensions

[deansher]

We execute the following statement if either rank is statically known, but it doesn't give the desired result if the other rank is unknown. https://github.com/tensorflow/java/blob/5532bf0af4a908ee36b6017a8745e515bb3bb544/tensorflow-framework/src/main/java/org/tensorflow/framework/losses/impl/LossesHelper.java#L215

Also, I believe we need to specify the last axis here: https://github.com/tensorflow/java/blob/5532bf0af4a908ee36b6017a8745e515bb3bb544/tensorflow-framework/src/main/java/org/tensorflow/framework/losses/impl/LossesHelper.java#L217 and here: https://github.com/tensorflow/java/blob/5532bf0af4a908ee36b6017a8745e515bb3bb544/tensorflow-framework/src/main/java/org/tensorflow/framework/losses/impl/LossesHelper.java#L219

[deansher]

and here: https://github.com/tensorflow/java/blob/5532bf0af4a908ee36b6017a8745e515bb3bb544/tensorflow-framework/src/main/java/org/tensorflow/framework/losses/impl/LossesHelper.java#L123

[adiroy05]

sir, i can help over it

[karllessard]

Please do @adiroy05 , feel free to submit a PR, thanks!