[Bug]: torch_where op raises NotImplementedError

System Info

H100

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

build_and_run_ad.py --model allenai/OLMoE-1B-7B-0924 --args.model-factory AutoModelForCausalLM '--args.model-kwargs={}' --args.tokenizer null --args.world-size 2 --args.compile-backend torch-compile --args.attn-backend flashinfer --args.runtime trtllm --args.skip-loading-weights False --args.transforms.detect-sharding.simple-shard-only False --args.max-seq-len 512 --benchmark.enabled True --benchmark.results-path /jet/logs/basic/auto-deploy-model-coverage_ab-flashinfer_b-true_cb-torch-compile_m-allenai-olmoe-1b-7b-0924_mf-automodelforcausallm_mk--_msl-512_r-trtllm_sso-false_sw-false_t-null_ws-2/extra.json --benchmark.store-results true

Expected behavior

should pass

actual behavior

0:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
0:     return forward_call(*args, **kwargs)
0:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0:   File "/usr/local/lib/python3.12/dist-packages/transformers/models/olmoe/modeling_olmoe.py", line 610, in forward
0:     idx, top_x = torch.where(expert_mask[expert_idx])
0:                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0:   File "/opt/tensorrt-llm/tensorrt_llm/_torch/auto_deploy/export/library/torch_where.py", line 25, in _torch_where_patch
0:     return torch.nonzero(condition, as_tuple=True)
0:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0:   File "/usr/local/lib/python3.12/dist-packages/torch/_prims_common/wrappers.py", line 309, in _fn
0:     result = fn(*args, **kwargs)
0:              ^^^^^^^^^^^^^^^^^^^
0:   File "/usr/local/lib/python3.12/dist-packages/torch/_meta_registrations.py", line 3332, in nonzero
0:     torch._check_not_implemented(
0:   File "/usr/local/lib/python3.12/dist-packages/torch/__init__.py", line 1786, in _check_not_implemented
0:     _check_with(NotImplementedError, cond, message)
0:   File "/usr/local/lib/python3.12/dist-packages/torch/__init__.py", line 1665, in _check_with
0:     raise error_type(message_evaluated)
0: NotImplementedError: The register_meta function for torch.nonzero() raises unimplemented by default, as a correct data-independent implementation does not exist. This implementation returns a fake value, assuming all elements of the tensor are non-zero. To enable this registration, please set 'torch.fx.experimental._config.meta_nonzero_assume_all_nonzero' to True.

additional notes

NA

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.