> by "the problem of bounding objects" you mean object detection / localization then a single regression head on top a classifier architecture is not a good way of solving this problem, there

I just replied myself a similar thing in a comment. You are correct,

Indeed I am planning to do it in Keras because there arent implemented models in TF.js and doing it is quite difficult.

I do not think Pytorch models can be easily used in Tensorflow JS afterwards right