fix nonzero upper bound (#393)

Summary:
Pull Request resolved: #393

If nonzero's input has `n` dimensions, then the resulting output tensor is of size `(z, n)`, where `z` is the total number of non-zero elements in the input tensor. Therefore, an upper bound for out's shape is [input.numel(), input.dim()]

Reviewed By: ydwu4

Differential Revision: D49342476

fbshipit-source-id: ac70fd83090cc80c93315b9fc5be830afb6207df

Author: manuelcandales

exir/passes/sym_shape_eval_pass.py

@@ -29,7 +29,7 @@ def inference_deco(f: Callable):
 @register_upper_bound_inference(exir_ops.edge.aten.nonzero.default)
 @register_upper_bound_inference(torch.ops.aten.nonzero.default)
 def nonzero(args, kwargs) -> List[Optional[int]]:
-    return [eval_expr(args[0].shape[0]), len(args[0].shape)]
+    return [eval_expr(args[0].numel()), len(args[0].shape)]
 
 
 class HintBasedSymShapeEvalPass(PassBase):