Hi @sksvineeth - thanks for this PR. It still requires much more work:

• first, the code is broken
• weights needs a default value
• there should be a consideration of what happens if a some values of weights is NaN
• there should be a consideration of what happens if the nan_strategy is set to drop, as some values a re trimmed.

Fixed he first two bullets, the remaining two requires more work

• first, the code is broken
• weights needs a default value
• there should be a consideration of what happens if a some values of weights is NaN
• there should be a consideration of what happens if the nan_strategy is set to drop, as some values a re trimmed

Closing PR as it is untouched for too long