HazyResearch/metal

metal/util/pred_to_prob Squeezing 1d tensor with 1 element to 0d tensor

dliangsta opened this issue · 1 comments

If Y_h is a 1d tensor with 1 element, and we squeeze that tensor, it becomes a 0d tensor, and fails the check in the subsequent line (56) at assert Y_h.dim() == 1

I'm not sure how to make a pull request because I can't seem to push a new branch to the repo (probably permissions) but I think the fix would be to check that Y_h.dim() > 1 before squeezing, because the two possible outcomes are a 0d tensor (bad) and the same 1d tensor (no effect), so we shouldnt squeeze when Y_h.dim() is 1.

Edit: I figured out how to make a pull request (with Jared's help)