Cherrypick fix for DLRM real dataset crash (pytorch#2409)

Co-authored-by: Davide Libenzi <[email protected]>

[r1.5] Cherrypick torch_xla/csrc/elementwise.cpp from (pytorch#1826)

pytorch@4d454fc