Repository: pytorch
Issue: In-place native functional collectives right after a custom torch op fail to codegen in inductor
Labels: oncall: distributed,triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

When compiling one of my models using the new native functional collectives I hit the following error:

```
[rank1]: Traceback (most recent call last):
[rank1]:   File "/lustre/aviros/fms-oss/scripts/inference.py", line 234, in <module>
[rank1]:     infer(cache, sample)
[rank1]:   File "/lustre/aviros/fms-oss/scripts/inference.py", line 215, in infer
[rank1]:     result = generate(
[rank1]:              ^^^^^^^^^
[rank1]:   File "/lustre/aviros/fms-oss/fms/utils/generation.py", line 111, in generate
[rank1]:     output = prefill_model(input_ids, **kwargs)
[rank1]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1536, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 437, in _fn
[rank1]:     return fn(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1536, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 921, in catch_errors
[rank1]:     return callback(frame, cache_entry, hooks, frame_state, skip=1)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert
[rank1]:     return _compile(
[rank1]:            ^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/contextlib.py", line 81, in inner
[rank1]:     return func(*args, **kwds)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile
[rank1]:     guarded_code = compile_inner(code, one_graph, hooks, transform)
[rank1]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank1]:     r = func(*args, **kwargs)
[rank1]:         ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner
[rank1]:     out_code = transform_code_object(code, transform)
[rank1]:                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/bytecode_transformation.py", line 1033, in transform_code_object
[rank1]:     transformations(instructions, code_options)
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn
[rank1]:     return fn(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform
[rank1]:     tracer.run()
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2195, in run
[rank1]:     super().run()
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 850, in run
[rank1]:     and self.step()
[rank1]:         ^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 813, in step
[rank1]:     getattr(self, inst.opname)(inst)
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2314, in RETURN_VALUE
[rank1]:     self.output.compile_subgraph(
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 984, in compile_subgraph
[rank1]:     self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/contextlib.py", line 81, in inner
[rank1]:     return func(*args, **kwds)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1161, in compile_and_call_fx_graph
[rank1]:     compiled_fn = self.call_user_compiler(gm)
[rank1]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank1]:     r = func(*args, **kwargs)
[rank1]:         ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1234, in call_user_compiler
[rank1]:     raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1215, in call_user_compiler
[rank1]:     compiled_fn = compiler_fn(gm, self.example_inputs())
[rank1]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/repro/after_dynamo.py", line 117, in debug_wrapper
[rank1]:     compiled_gm = compiler_fn(gm, example_inputs)
[rank1]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/repro/after_dynamo.py", line 117, in debug_wrapper
[rank1]:     compiled_gm = compiler_fn(gm, example_inputs)
[rank1]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/__init__.py", line 1729, in __call__
[rank1]:     return compile_fx(model_, inputs_, config_patches=self.config)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/contextlib.py", line 81, in inner
[rank1]:     return func(*args, **kwds)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 1373, in compile_fx
[rank1]:     return aot_autograd(
[rank1]:            ^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/backends/common.py", line 58, in compiler_fn
[rank1]:     cg = aot_module_simplified(gm, example_inputs, **kwargs)
[rank1]:          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 878, in aot_module_simplified
[rank1]:     compiled_fn = create_aot_dispatcher_function(
[rank1]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank1]:     r = func(*args, **kwargs)
[rank1]:         ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 603, in create_aot_dispatcher_function
[rank1]:     compiled_fn = compiler_fn(flat_fn, fake_flat_args, aot_config, fw_metadata=fw_metadata)
[rank1]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 434, in aot_wrapper_dedupe
[rank1]:     return compiler_fn(flat_fn, leaf_flat_args, aot_config, fw_metadata=fw_metadata)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 639, in aot_wrapper_synthetic_base
[rank1]:     return compiler_fn(flat_fn, flat_args, aot_config, fw_metadata=fw_metadata)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 119, in aot_dispatch_base
[rank1]:     compiled_fw = compiler(fw_module, updated_flat_args)
[rank1]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank1]:     r = func(*args, **kwargs)
[rank1]:         ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 1300, in fw_compiler_base
[rank1]:     return inner_compile(
[rank1]:            ^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/repro/after_aot.py", line 83, in debug_wrapper
[rank1]:     inner_compiled_fn = compiler_fn(gm, example_inputs)
[rank1]:                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/debug.py", line 304, in inner
[rank1]:     return fn(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/contextlib.py", line 81, in inner
[rank1]:     return func(*args, **kwds)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/contextlib.py", line 81, in inner
[rank1]:     return func(*args, **kwds)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank1]:     r = func(*args, **kwargs)
[rank1]:         ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 463, in compile_fx_inner
[rank1]:     compiled_graph = fx_codegen_and_compile(
[rank1]:                      ^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 738, in fx_codegen_and_compile
[rank1]:     compiled_fn = graph.compile_to_fn()
[rank1]:                   ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/graph.py", line 1295, in compile_to_fn
[rank1]:     return self.compile_to_module().call
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank1]:     r = func(*args, **kwargs)
[rank1]:         ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/graph.py", line 1238, in compile_to_module
[rank1]:     self.codegen_with_cpp_wrapper() if self.cpp_wrapper else self.codegen()
[rank1]:                                                              ^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/graph.py", line 1195, in codegen
[rank1]:     self.scheduler = Scheduler(self.buffers)
[rank1]:                      ^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank1]:     r = func(*args, **kwargs)
[rank1]:         ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/scheduler.py", line 1260, in __init__
[rank1]:     self.nodes = [self.create_scheduler_node(n) for n in nodes]
[rank1]:                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/scheduler.py", line 1260, in <listcomp>
[rank1]:     self.nodes = [self.create_scheduler_node(n) for n in nodes]
[rank1]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/scheduler.py", line 1354, in create_scheduler_node
[rank1]:     return ExternKernelSchedulerNode(self, node)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/scheduler.py", line 166, in __init__
[rank1]:     self.set_read_writes(node.get_read_writes())
[rank1]:                          ^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/ir.py", line 7879, in get_read_writes
[rank1]:     volatile_reads = self.get_volatile_reads()
[rank1]:                      ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/lustre/aviros/conda_envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/ir.py", line 7848, in get_volatile_reads
[rank1]:     assert isinstance(coll, _CollectiveKernel)
[rank1]: torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
[rank1]: AssertionError: 
```

When investigating why this is the case, I read through the `get_volatile_reads` function, and I noticed that the function expects that `inp` being a `MultiOutput` means that the collective is out-of-place. When running an in-place collective like all-reduce, whose input comes from a custom pytorch op (aka defined though `torch.library`), `inp` will also be a `MultiOutput` (as that's what a `FallbackKernel` will generate), but the collective will still be in-place. A pretty easy and self-contained fix I found is changing the `get_volatile_reads` function to the following:

```
def get_volatile_reads_fixed(self):
    inp = self.inputs[0]
    if isinstance(inp, ir._CollectiveKernel):
        # Out-of-place single-output
        return [inp.inputs[0]]
    elif isinstance(inp, ir.MultiOutput):
        # Out-of-place multi-output
        coll = inp.inputs[0]
        if isinstance(coll, ir._CollectiveKernel):
            _, idx = inp.indices[0]
            return [coll.inputs[idx]]
        return []  # e.g. regular FallbackKernel
    else:
        # In-place requires no additional deps handling for volatile
        # reads since the inputs are mutated.
        return []
```

I don't know if this is semantically correct, but I can make a PR if you think this is the correct fix.

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @msaroufim @bdhirsh @anijain2305 @zou3519 @peterbell10 @aakhundov @yifuwang 

### Versions

Collecting environment information...
PyTorch version: 2.3.0.dev20240304+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.11.6 | packaged by conda-forge | (main, Oct  3 2023, 10:40:35) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-1041-aws-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 535.54.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   46 bits physical, 48 bits virtual
CPU(s):                          96
On-line CPU(s) list:             0-95
Thread(s) per core:              2
Core(s) per socket:              24
Socket(s):                       2
NUMA node(s):                    2
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           85
Model name:                      Intel(R) Xeon(R) Platinum 8275CL CPU @ 3.00GHz
Stepping:                        7
CPU MHz:                         3599.684
BogoMIPS:                        5999.99
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       1.5 MiB
L1i cache:                       1.5 MiB
L2 cache:                        48 MiB
L3 cache:                        71.5 MiB
NUMA node0 CPU(s):               0-23,48-71
NUMA node1 CPU(s):               24-47,72-95
Vulnerability Itlb multihit:     KVM: Mitigation: VMX unsupported
Vulnerability L1tf:              Mitigation; PTE Inversion
Vulnerability Mds:               Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Meltdown:          Mitigation; PTI
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:          Vulnerable
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves ida arat pku ospke

Versions of relevant libraries:
[pip3] mypy==1.7.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.2
[pip3] pytorch-triton==3.0.0+901819d2b6
[pip3] torch==2.3.0.dev20240304+cu121
[pip3] torchaudio==2.2.0.dev20240304+cu121
[pip3] torchvision==0.18.0.dev20240304+cu121
[conda] libblas                   3.9.0            16_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            16_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            16_linux64_mkl    conda-forge
[conda] mkl                       2022.1.0           h84fe81f_915    conda-forge
[conda] numexpr                   2.8.3           mkl_py311h0b1cd97_1    conda-forge
[conda] numpy                     1.26.2          py311h64a7726_0    conda-forge
[conda] pytorch-cuda              12.1                 ha16c6d3_5    pytorch
[conda] pytorch-triton            3.0.0+901819d2b6          pypi_0    pypi
[conda] torch                     2.3.0.dev20240304+cu121          pypi_0    pypi
[conda] torchaudio                2.2.0.dev20240304+cu121          pypi_0    pypi
[conda] torchvision               0.18.0.dev20240304+cu121          pypi_0    pypi