Repository: pytorch
Issue: nonzero with dtensor hang
Labels: oncall: distributed,triaged

### 🐛 Describe the bug

```
#0  0x00007fb5a545ae50 in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#1  0x00007fb5a51aa479 in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#2  0x00007fb5a553d95f in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#3  0x00007fb5a553e54e in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#4  0x00007fb5a51bc74d in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#5  0x00007fb5a53fbf3c in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#6  0x00007fb5a53fc831 in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#7  0x00007fb5a55353ab in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#8  0x00007fb5a5174f66 in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#9  0x00007fb5a5175379 in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#10 0x00007fb5a51761b7 in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#11 0x00007fb5a5363fa1 in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#12 0x00007fb6123f4459 in ?? () from /usr/local/cuda/lib64/libcudart.so.11.0
#13 0x00007fb6123c76fd in ?? () from /usr/local/cuda/lib64/libcudart.so.11.0
#14 0x00007fb6124196a5 in cudaMemcpyAsync () from /usr/local/cuda/lib64/libcudart.so.11.0
#15 0x00007fb5fbd62480 in c10::cuda::memcpy_and_sync (stream=0x0, kind=cudaMemcpyDeviceToHost, nbytes=4, src=0x7fb562a6de00, dst=0x7ffe9129b5e0) at /root/cery.69/pytorch/c10/cuda/CUDAFunctions.h:96
#16 at::native::nonzero_cuda_out_impl<long> (self=..., out=...) at /root/cery.69/pytorch/aten/src/ATen/native/cuda/Nonzero.cu:79
#17 0x00007fb5fbd553aa in operator() (__closure=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/native/cuda/Nonzero.cu:131
#18 operator() (__closure=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/native/cuda/Nonzero.cu:131
#19 at::native::nonzero_out_cuda (self=..., out=...) at /root/cery.69/pytorch/aten/src/ATen/native/cuda/Nonzero.cu:131
#20 0x00007fb5fbd5561d in at::native::nonzero_cuda (self=...) at /root/cery.69/pytorch/aten/src/ATen/native/cuda/Nonzero.cu:139
#21 0x00007fb5fc9ddc2f in at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA__nonzero (self=...) at /root/cery.69/pytorch/build/aten/src/ATen/RegisterCUDA.cpp:23876
#22 0x00007fb5fc9ddcad in c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor(const at::Tensor&), at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA__nonzero>, at::Tensor, c10::guts::typelist::typelist<const at::Tensor&> >::operator() (args#0=..., this=<optimized out>)
    at /root/cery.69/pytorch/aten/src/ATen/core/boxing/impl/WrapFunctionIntoFunctor.h:13
#23 c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor(const at::Tensor&), at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA__nonzero>, at::Tensor, c10::guts::typelist::typelist<const at::Tensor&> >, at::Tensor(const at::Tensor&)>::call(c10::OperatorKernel *, c10::DispatchKeySet, const at::Tensor &) (functor=<optimized out>, args#0=...) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h:468
#24 0x00007fb608041e4f in c10::callUnboxedKernelFunction<at::Tensor, at::Tensor const&> (dispatchKeySet=..., functor=<optimized out>, unboxed_kernel_func=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/KernelFunction_impl.h:47
#25 c10::KernelFunction::call<at::Tensor, at::Tensor const&> (dispatchKeySet=..., opHandle=..., this=0x1da7668) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/KernelFunction_impl.h:103
#26 c10::Dispatcher::call<at::Tensor, at::Tensor const&>(c10::TypedOperatorHandle<at::Tensor (at::Tensor const&)> const&, at::Tensor const&) const (op=..., this=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/core/dispatch/Dispatcher.h:673
#27 c10::TypedOperatorHandle<at::Tensor (at::Tensor const&)>::call(at::Tensor const&) const (args#0=..., this=0x7fb610ea8cd0 <at::_ops::nonzero::call(at::Tensor const&)::op>) at /root/cery.69/pytorch/aten/src/ATen/core/dispatch/Dispatcher.h:521
#28 at::_ops::nonzero::call (self=...) at /root/cery.69/pytorch/build/aten/src/ATen/Operators_0.cpp:9651
#29 0x00007fb61166fc23 in at::Tensor::nonzero (this=0x7ffe9129bc68) at /root/cery.69/pytorch/build/aten/src/ATen/core/TensorBody.h:5041
#30 torch::autograd::dispatch_nonzero (self=...) at /root/cery.69/pytorch/torch/csrc/autograd/python_torch_functions_manual.cpp:153
#31 torch::autograd::THPVariable_nonzero (self=<optimized out>, args=0x7fb56c5e7280, kwargs=0x7ffe9129bc68) at /root/cery.69/pytorch/torch/csrc/autograd/python_torch_functions_manual.cpp:903
#32 0x000000000053f350 in ?? ()
#33 0x000000000051d89b in _PyObject_MakeTpCall ()
#34 0x00000000005175ba in _PyEval_EvalFrameDefault ()
#35 0x0000000000528b63 in _PyFunction_Vectorcall ()
#36 0x0000000000511fb5 in _PyEval_EvalFrameDefault ()
#37 0x0000000000528b63 in _PyFunction_Vectorcall ()
#38 0x0000000000512192 in _PyEval_EvalFrameDefault ()
#39 0x0000000000528b63 in _PyFunction_Vectorcall ()
#40 0x0000000000511fb5 in _PyEval_EvalFrameDefault ()
#41 0x0000000000528b63 in _PyFunction_Vectorcall ()
#42 0x0000000000512192 in _PyEval_EvalFrameDefault ()
#43 0x0000000000510fe7 in ?? ()
#44 0x0000000000528d21 in _PyFunction_Vectorcall ()
#45 0x0000000000513e8b in _PyEval_EvalFrameDefault ()
#46 0x00000000005106ed in ?? ()
#47 0x0000000000528d21 in _PyFunction_Vectorcall ()
#48 0x0000000000511fb5 in _PyEval_EvalFrameDefault ()
#49 0x0000000000510fe7 in ?? ()
#50 0x0000000000528d21 in _PyFunction_Vectorcall ()
#51 0x0000000000516e76 in _PyEval_EvalFrameDefault ()
#52 0x00000000005106ed in ?? ()
#53 0x0000000000528d21 in _PyFunction_Vectorcall ()
#54 0x000000000053bdb1 in ?? ()
#55 0x000000000052842e in ?? ()
#56 0x00000000005282d8 in PyObject_CallFunctionObjArgs ()
#57 0x00007fb6119ea79e in torch::dispatch_on_subclass (maybe_mode_key=..., torch_function_name_str=0x7fb611b87948 "__torch_dispatch__", is_torch_function=<optimized out>, torch_api_function=0x7fb57fca6460, py_types=..., overloaded_args=..., kwargs=0x7fb56c5ef100, args=0x7fb56c63cfd0)
    at /root/cery.69/pytorch/torch/include/pybind11/pytypes.h:242
#58 torch::handle_torch_function_no_python_arg_parser (overloaded_args=..., args=args@entry=0x7fb56c63cfd0, kwargs=kwargs@entry=0x7fb56c5ef100, func_name=func_name@entry=0x0, torch_api_function=torch_api_function@entry=0x7fb57fca6460, module_name=<optimized out>, torch_function_name=<optimized out>)
    at /root/cery.69/pytorch/torch/csrc/utils/python_arg_parser.cpp:495
#59 0x00007fb6115c8408 in (anonymous namespace)::ConcretePyInterpreterVTable::dispatch (this=<optimized out>, op=..., stack=0x7ffe9129da30) at /usr/include/c++/10/bits/stl_vector.h:918
#60 0x00007fb6076ca939 in (anonymous namespace)::pythonFallback (op=..., stack=0x7ffe9129da30) at /root/cery.69/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:87
#61 0x00007fb60aec32c5 in c10::BoxedKernel::callBoxed (stack=0x7ffe9129da30, dispatchKeySet=..., opHandle=..., this=0x1da7ee8) at /root/cery.69/pytorch/c10/util/intrusive_ptr.h:403
#62 c10::KernelFunction::callBoxed (stack=0x7ffe9129da30, dispatchKeySet=..., opHandle=..., this=0x1da7ee8) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/KernelFunction_impl.h:43
#63 c10::Dispatcher::callBoxed (op=..., stack=0x7ffe9129da30, this=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/core/dispatch/Dispatcher.h:725
#64 0x00007fb6117c8b53 in std::function<void (std::vector<c10::IValue, std::allocator<c10::IValue> >&)>::operator()(std::vector<c10::IValue, std::allocator<c10::IValue> >&) const (__args#0=std::vector of length 0, capacity 1, this=0x7ffe9129da50) at /usr/include/c++/10/bits/std_function.h:622
#65 torch::jit::Operation::operator() (stack=std::vector of length 0, capacity 1, this=0x7ffe9129da50) at /root/cery.69/pytorch/aten/src/ATen/core/stack.h:41
#66 torch::jit::invokeOperatorFromPython (operations=..., args=..., kwargs=..., dk=std::optional<c10::DispatchKey> [no contained value]) at /root/cery.69/pytorch/torch/csrc/jit/python/pybind_utils.cpp:773
#67 0x00007fb6117c9221 in torch::jit::_get_operation_for_overload_or_packet (operations=std::vector of length 1, capacity 1 = {...}, symbol=..., args=..., kwargs=..., is_overload=is_overload@entry=true, dk=std::optional<c10::DispatchKey> [no contained value]) at /root/cery.69/pytorch/torch/include/pybind11/pytypes.h:259
#68 0x00007fb6116d186a in operator() (__closure=0x6055fb8, args=..., kwargs=...) at /root/cery.69/pytorch/torch/csrc/jit/python/init.cpp:1604
#69 0x00007fb6116d1dcc in pybind11::detail::argument_loader<pybind11::args, pybind11::kwargs>::call_impl<pybind11::object, torch::jit::initJITBindings(PyObject*)::<lambda(const string&, const string&)>::<lambda(pybind11::args, pybind11::kwargs)>&, 0, 1, pybind11::detail::void_type> (f=..., this=0x7ffe9129dce0)
    at /root/cery.69/pytorch/torch/include/pybind11/pytypes.h:359
#70 pybind11::detail::argument_loader<pybind11::args, pybind11::kwargs>::call<pybind11::object, pybind11::detail::void_type, torch::jit::initJITBindings(PyObject*)::<lambda(const string&, const string&)>::<lambda(pybind11::args, pybind11::kwargs)>&> (f=..., this=0x7ffe9129dce0) at /root/cery.69/pytorch/torch/include/pybind11/cast.h:1449
#71 operator() (this=0x0, call=...) at /root/cery.69/pytorch/torch/include/pybind11/pybind11.h:254
#72 _FUN () at /root/cery.69/pytorch/torch/include/pybind11/pybind11.h:224
#73 0x00007fb6112d2cf9 in pybind11::cpp_function::dispatcher (self=<optimized out>, args_in=0x7fb56c5e7190, kwargs_in=0x7fb56c6340c0) at /root/cery.69/pytorch/torch/include/pybind11/pybind11.h:946
#74 0x000000000053f350 in ?? ()
#75 0x000000000053c3c9 in PyObject_Call ()
#76 0x0000000000517b9b in _PyEval_EvalFrameDefault ()
#77 0x00000000005106ed in ?? ()
#78 0x0000000000528d21 in _PyFunction_Vectorcall ()
#79 0x000000000051cf95 in _PyObject_FastCallDictTstate ()
#80 0x0000000000538545 in _PyObject_Call_Prepend ()
#81 0x00000000005b81f3 in ?? ()
#82 0x000000000053c3c9 in PyObject_Call ()
#83 0x0000000000513e8b in _PyEval_EvalFrameDefault ()
#84 0x00000000005106ed in ?? ()
#85 0x0000000000528d21 in _PyFunction_Vectorcall ()
#86 0x000000000053bdb1 in ?? ()
#87 0x00000000005302b5 in ?? ()
#88 0x00000000005f702d in PyObject_CallMethod ()
#89 0x00007fb6119eb0d3 in torch::dispatch_on_mode (torch_function_name_str=0x7fb611b87948 "__torch_dispatch__", is_torch_function=<optimized out>, torch_api_function=<optimized out>, py_types=..., kwargs=<optimized out>, args=<optimized out>) at /root/cery.69/pytorch/torch/include/pybind11/pytypes.h:242
#90 torch::handle_torch_function_no_python_arg_parser (overloaded_args=..., args=args@entry=0x7fb5740bac40, kwargs=kwargs@entry=0x7fb56c634580, func_name=func_name@entry=0x0, torch_api_function=torch_api_function@entry=0x7fb57fca6460, module_name=<optimized out>, torch_function_name=<optimized out>)
    at /root/cery.69/pytorch/torch/csrc/utils/python_arg_parser.cpp:476
#91 0x00007fb6115c8408 in (anonymous namespace)::ConcretePyInterpreterVTable::dispatch (this=<optimized out>, op=..., stack=0x7ffe9129f000) at /usr/include/c++/10/bits/stl_vector.h:918
#92 0x00007fb6076ca939 in (anonymous namespace)::pythonFallback (op=..., stack=0x7ffe9129f000) at /root/cery.69/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:87
#93 0x00007fb6080d9521 in c10::BoxedKernel::callBoxed (stack=0x7ffe9129f000, dispatchKeySet=..., opHandle=..., this=0x1da7ee8) at /root/cery.69/pytorch/c10/util/intrusive_ptr.h:403
#94 c10::impl::BoxedKernelWrapper<at::Tensor (at::Tensor const&), void>::call(c10::BoxedKernel const&, c10::OperatorHandle const&, c10::DispatchKeySet, at::Tensor const&) (boxed_kernel_func=..., opHandle=..., dispatchKeySet=dispatchKeySet@entry=..., args#0=...) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/impl/boxing.h:227
#95 0x00007fb607fecd68 in c10::KernelFunction::call<at::Tensor, at::Tensor const&> (dispatchKeySet=..., opHandle=..., this=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/KernelFunction_impl.h:112
#96 c10::Dispatcher::redispatch<at::Tensor, at::Tensor const&>(c10::TypedOperatorHandle<at::Tensor (at::Tensor const&)> const&, c10::DispatchKeySet, at::Tensor const&) const (this=<optimized out>, currentDispatchKeySet=..., op=...) at /root/cery.69/pytorch/aten/src/ATen/core/dispatch/Dispatcher.h:690
#97 c10::TypedOperatorHandle<at::Tensor (at::Tensor const&)>::redispatch(c10::DispatchKeySet, at::Tensor const&) const (args#0=..., currentDispatchKeySet=..., this=0x7fb610ea8cb0 <at::_ops::nonzero::redispatch(c10::DispatchKeySet, at::Tensor const&)::op>) at /root/cery.69/pytorch/aten/src/ATen/core/dispatch/Dispatcher.h:526
#98 at::_ops::nonzero::redispatch (dispatchKeySet=..., self=...) at /root/cery.69/pytorch/build/aten/src/ATen/Operators_0.cpp:9658
#99 0x00007fb609cc0fd2 in at::redispatch::nonzero (self=..., dispatchKeySet=...) at /root/cery.69/pytorch/build/aten/src/ATen/RedispatchFunctions.h:10797
#100 operator() (__closure=<optimized out>) at /root/cery.69/pytorch/torch/csrc/autograd/generated/VariableType_0.cpp:13214
#101 torch::autograd::VariableType::(anonymous namespace)::nonzero (ks=..., self=...) at /root/cery.69/pytorch/torch/csrc/autograd/generated/VariableType_0.cpp:13215
#102 0x00007fb609cd66a4 in c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor(c10::DispatchKeySet, const at::Tensor&), torch::autograd::VariableType::(anonymous namespace)::nonzero>, at::Tensor, c10::guts::typelist::typelist<c10::DispatchKeySet, const at::Tensor&> >::operator() (args#1=..., args#0=...,
    this=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/impl/WrapFunctionIntoFunctor.h:13
#103 c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor(c10::DispatchKeySet, const at::Tensor&), torch::autograd::VariableType::(anonymous namespace)::nonzero>, at::Tensor, c10::guts::typelist::typelist<c10::DispatchKeySet, const at::Tensor&> >, at::Tensor(c10::DispatchKeySet, const at::Tensor&)>::call (args#0=..., dispatchKeySet=..., functor=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h:485
#104 c10::impl::call_functor_with_args_from_stack_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor(c10::DispatchKeySet, const at::Tensor&), torch::autograd::VariableType::(anonymous namespace)::nonzero>, at::Tensor, c10::guts::typelist::typelist<c10::DispatchKeySet, const at::Tensor&> >, false, 0, const at::Tensor&> (stack=0x7ffe9129f170, dispatchKeySet=..., functor=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h:506
#105 c10::impl::call_functor_with_args_from_stack<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor(c10::DispatchKeySet, const at::Tensor&), torch::autograd::VariableType::(anonymous namespace)::nonzero>, at::Tensor, c10::guts::typelist::typelist<c10::DispatchKeySet, const at::Tensor&> >, false> (
    stack=0x7ffe9129f170, dispatchKeySet=..., functor=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h:518
#106 c10::impl::make_boxed_from_unboxed_functor<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor(c10::DispatchKeySet, const at::Tensor&), torch::autograd::VariableType::(anonymous namespace)::nonzero>, at::Tensor, c10::guts::typelist::typelist<c10::DispatchKeySet, const at::Tensor&> >, false>::call(c10::OperatorKernel *, const c10::OperatorHandle &, c10::DispatchKeySet, c10::Stack *) (functor=<optimized out>, dispatchKeySet=..., stack=0x7ffe9129f170) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h:584
#107 0x00007fb6076ca7ec in (anonymous namespace)::pythonTLSSnapshotFallback (stack=0x7ffe9129f170, dispatch_keys=..., op=...) at /root/cery.69/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:108
#108 c10::BoxedKernel::make_boxed_function<(anonymous namespace)::pythonTLSSnapshotFallback> (opHandle=..., ks=..., stack=0x7ffe9129f170) at /root/cery.69/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:163
#109 0x00007fb6080d9521 in c10::BoxedKernel::callBoxed (stack=0x7ffe9129f170, dispatchKeySet=..., opHandle=..., this=0x1da8408) at /root/cery.69/pytorch/c10/util/intrusive_ptr.h:403
#110 c10::impl::BoxedKernelWrapper<at::Tensor (at::Tensor const&), void>::call(c10::BoxedKernel const&, c10::OperatorHandle const&, c10::DispatchKeySet, at::Tensor const&) (boxed_kernel_func=..., opHandle=..., dispatchKeySet=dispatchKeySet@entry=..., args#0=...) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/impl/boxing.h:227
#111 0x00007fb608041fa0 in c10::KernelFunction::call<at::Tensor, at::Tensor const&> (dispatchKeySet=..., opHandle=..., this=0x1da8408) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/KernelFunction_impl.h:112
#112 c10::Dispatcher::call<at::Tensor, at::Tensor const&>(c10::TypedOperatorHandle<at::Tensor (at::Tensor const&)> const&, at::Tensor const&) const (op=..., this=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/core/dispatch/Dispatcher.h:673
#113 c10::TypedOperatorHandle<at::Tensor (at::Tensor const&)>::call(at::Tensor const&) const (args#0=..., this=0x7fb610ea8cd0 <at::_ops::nonzero::call(at::Tensor const&)::op>) at /root/cery.69/pytorch/aten/src/ATen/core/dispatch/Dispatcher.h:521
#114 at::_ops::nonzero::call (self=...) at /root/cery.69/pytorch/build/aten/src/ATen/Operators_0.cpp:9651
#115 0x00007fb607ae0e5d in at::Tensor::nonzero (this=<optimized out>) at /root/cery.69/pytorch/build/aten/src/ATen/core/TensorBody.h:5041
#116 at::native::nonzero_numpy (self=...) at /root/cery.69/pytorch/aten/src/ATen/native/TensorAdvancedIndexing.cpp:2459
#117 0x00007fb608a4bf4a in at::(anonymous namespace)::(anonymous namespace)::wrapper_CompositeImplicitAutograd__nonzero_numpy (self=...) at /root/cery.69/pytorch/build/aten/src/ATen/RegisterCompositeImplicitAutograd.cpp:5033
#118 c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<std::vector<at::Tensor, std::allocator<at::Tensor> >(const at::Tensor&), at::(anonymous namespace)::(anonymous namespace)::wrapper_CompositeImplicitAutograd__nonzero_numpy>, std::vector<at::Tensor, std::allocator<at::Tensor> >, c10::guts::typelist::typelist<cons
t at::Tensor&> >::operator() (args#0=..., this=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/impl/WrapFunctionIntoFunctor.h:13
#119 c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<std::vector<at::Tensor, std::allocator<at::Tensor> >(const at::Tensor&), at::(anonymous namespace)::(anonymous namespace)::wrapper_CompositeImplicitAutograd__nonzero_numpy>, std::vector<at::Tensor, std::allocator<at::Tensor> >, c10::guts::typelist::typelist<const at::Tensor&> >, std::vector<at::Tensor, std::allocator<at::Tensor> >(const at::Tensor&)>::call (args#0=..., functor=<optimized out>) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h:468
#120 c10::impl::call_functor_with_args_from_stack_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<std::vector<at::Tensor>(const at::Tensor&), at::(anonymous namespace)::(anonymous namespace)::wrapper_CompositeImplicitAutograd__nonzero_numpy>, std::vector<at::Tensor>, c10::guts::typelist::typelist<const at::Tensor&> >, false, 0, const at::Tensor&> (functor=<optimized out>, dispatchKeySet=..., stack=0x7ffe9129f360) at /root/cery.69/pytorch/aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h:506
#121 c10::impl::call_functor_with_args_from_stack<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<std::vector<at::Tensor>(const at::Tensor&), at::(anonymous namespace)::(anonymous namespace)::wrapper_CompositeImplicitAutograd__nonzero_numpy>, std::vector<at::Tensor>, c10::guts::typelist::typelist<const at::Tensor&> >, false> (stack=0x7ffe9129f360, dis
```

i call nonzero with DTensor. in torch run mode it hang 

### Versions

torch 2.2.1-rc3

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang