Could not run select_backward [vmap] [dlrm] [functorch]
Labels: triaged, module: functorch

### 🐛 Describe the bug

What I am trying to do:

1. Create an input and target tensor
2. Create a list of models for ensemble (dlrm models from torchbench)
3. Use vmap (from functorch) to create an ensemble
4. Obtain the vmap function and pass it to aot_function (aot_autograd), to obtain the forward and backward graphs

NOTE: This workflow has been tested for hf_Bert, resnet and GPT2, and works fine for them.
Repro:

```
import torch
import torch.fx as fx
from functorch import (combine_state_for_ensemble,
                       make_functional_with_buffers, vmap)
from functorch._src.named_members_polyfill import (_named_buffers,
                                                   _named_parameters)
from functorch.compile import aot_function, aot_module, config, print_compile
from torch._subclasses import FakeTensor, FakeTensorMode

config.use_fake_tensor = True
from torchbenchmark import load_model_by_name


def fake_compiler(fx_g: fx.GraphModule, inps):
    print(fx_g.code)
    output_node = [node for node in fx_g.graph.nodes if node.op == 'output'][0]
    output_data = [node.meta['val'] if node is not None else node for node in output_node.args[0]]
    def new_f(*args):
        return output_data
    return new_f



Model = load_model_by_name("dlrm")
num_models = 2
batch_size = 1000

model_list = [
    Model(device="cuda", test="train", batch_size = batch_size) for _ in range(num_models)
]
b_models = [model_list[i].model for i in range(num_models)]

inp = model_list[0].example_inputs
loss_fn = model_list[0].loss_fn
targets = model_list[0].targets

print(torch.cuda.memory_allocated())

func_model, params, buffers = combine_state_for_ensemble(b_models)
for p in params:
    p.requires_grad = True



def compute_loss_dlrm(params, buffers, batch, targets):
    gen = func_model(params, buffers, *batch)
    loss = loss_fn(gen, targets)
    return loss


parallel_func = vmap(compute_loss_dlrm, in_dims=(0, 0, None, None), randomness="same")

aot_func = aot_function(parallel_func, fake_compiler)
out = aot_func(params, buffers, inp, targets)
print(out.size())
print(type(out))
print(out)
print(out.device)
out.sum().backward()
```


Error: Stack Trace

```
/scratch/sanketpurandare/work/pytorch/torch/nn/functional.py:2388: UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::embedding_bag.padding_idx. Please file us an issue on GitHub so that we can prioritize its implementation. (Triggered internally at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/functorch/BatchedFallback.cpp:84.)
  ret, _, _, _ = torch.embedding_bag(
/scratch/sanketpurandare/work/pytorch/torch/nn/functional.py:2388: UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::embedding_bag.padding_idx. Please file us an issue on GitHub so that we can prioritize its implementation. (Triggered internally at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/functorch/BatchedFallback.cpp:84.)
  ret, _, _, _ = torch.embedding_bag(
Traceback (most recent call last):
  File "/scratch/sanketpurandare/work/pytorch/torch/_subclasses/fake_tensor.py", line 821, in __torch_dispatch__
    r = func(*args, **kwargs)
  File "/scratch/sanketpurandare/work/pytorch/torch/_ops.py", line 257, in __call__
    return self._op(*args, **kwargs or {})
NotImplementedError: Could not run 'aten::select_backward' with arguments from the 'SparseMeta' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'aten::select_backward' is only available for these backends: [CPU, CUDA, HIP, MPS, IPU, XPU, HPU, VE, Meta, PrivateUse1, PrivateUse2, PrivateUse3, FPGA, ORT, Vulkan, Metal, QuantizedCPU, QuantizedCUDA, QuantizedHIP, QuantizedMPS, QuantizedIPU, QuantizedXPU, QuantizedHPU, QuantizedVE, QuantizedMeta, QuantizedPrivateUse1, QuantizedPrivateUse2, QuantizedPrivateUse3, CustomRNGKeyId, MkldnnCPU, SparseCsrCPU, SparseCsrCUDA, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradHIP, AutogradXLA, AutogradMPS, AutogradIPU, AutogradXPU, AutogradHPU, AutogradVE, AutogradLazy, AutogradMeta, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, AutogradNestedTensor, Tracer, AutocastCPU, AutocastCUDA, FuncTorchBatched, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PythonDispatcher].

Undefined: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
CPU: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
CUDA: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
HIP: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
MPS: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
IPU: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
XPU: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
HPU: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
VE: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
Meta: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
PrivateUse1: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
PrivateUse2: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
PrivateUse3: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
FPGA: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
ORT: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
Vulkan: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
Metal: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
QuantizedCPU: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
QuantizedCUDA: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
QuantizedHIP: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
QuantizedMPS: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
QuantizedIPU: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
QuantizedXPU: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
QuantizedHPU: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
QuantizedVE: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
QuantizedMeta: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
QuantizedPrivateUse1: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
QuantizedPrivateUse2: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
QuantizedPrivateUse3: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
CustomRNGKeyId: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
MkldnnCPU: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
SparseCsrCPU: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
SparseCsrCUDA: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
BackendSelect: fallthrough registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/core/BackendSelectFallbackKernel.cpp:3 [backend fallback]
Python: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:140 [backend fallback]
FuncTorchDynamicLayerBackMode: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/functorch/DynamicLayer.cpp:488 [backend fallback]
Functionalize: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/FunctionalizeFallbackKernel.cpp:291 [backend fallback]
Named: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/core/NamedRegistrations.cpp:7 [backend fallback]
Conjugate: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/ConjugateFallback.cpp:18 [backend fallback]
Negative: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/native/NegateFallback.cpp:18 [backend fallback]
ZeroTensor: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/ZeroTensorFallback.cpp:86 [backend fallback]
ADInplaceOrView: fallthrough registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:64 [backend fallback]
AutogradOther: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradCPU: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradCUDA: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradHIP: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradXLA: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradMPS: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradIPU: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradXPU: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradHPU: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradVE: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradLazy: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradMeta: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradPrivateUse1: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradPrivateUse2: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradPrivateUse3: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
AutogradNestedTensor: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/VariableType_1.cpp:14589 [autograd kernel]
Tracer: registered at /scratch/sanketpurandare/work/pytorch/torch/csrc/autograd/generated/TraceType_1.cpp:15586 [kernel]
AutocastCPU: fallthrough registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/autocast_mode.cpp:460 [backend fallback]
AutocastCUDA: fallthrough registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/autocast_mode.cpp:331 [backend fallback]
FuncTorchBatched: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/functorch/BatchRulesViews.cpp:512 [kernel]
FuncTorchVmapMode: fallthrough registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/functorch/VmapModeRegistrations.cpp:28 [backend fallback]
Batched: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/BatchingRegistrations.cpp:1068 [kernel]
VmapMode: fallthrough registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/VmapModeRegistrations.cpp:33 [backend fallback]
FuncTorchGradWrapper: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/functorch/TensorWrapper.cpp:189 [backend fallback]
PythonTLSSnapshot: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:148 [backend fallback]
FuncTorchDynamicLayerFrontMode: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/functorch/DynamicLayer.cpp:484 [backend fallback]
PythonDispatcher: registered at /scratch/sanketpurandare/work/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:144 [backend fallback]


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/scratch/sanketpurandare/work/testing/test_patch.py", line 59, in <module>
    out = aot_func(params, buffers, inp, targets)
  File "/scratch/sanketpurandare/work/pytorch/functorch/_src/aot_autograd.py", line 671, in returned_function
    compiled_fn = create_aot_dispatcher_function(
  File "/scratch/sanketpurandare/work/pytorch/functorch/_src/aot_autograd.py", line 524, in create_aot_dispatcher_function
    aot_dispatch_autograd(flat_fn, fake_flat_tensor_args, aot_config)
  File "/scratch/sanketpurandare/work/pytorch/functorch/_src/aot_autograd.py", line 367, in aot_dispatch_autograd
    fx_g = make_fx(joint_forward_backward)(*joint_inputs)
  File "/scratch/sanketpurandare/work/pytorch/torch/fx/experimental/proxy_tensor.py", line 660, in wrapped
    t = dispatch_trace(wrap_key(func, args, fx_tracer), tracer=fx_tracer, concrete_args=tuple(phs))
  File "/scratch/sanketpurandare/work/pytorch/torch/fx/experimental/proxy_tensor.py", line 408, in dispatch_trace
    graph = tracer.trace(root, concrete_args)
  File "/scratch/sanketpurandare/work/pytorch/torch/fx/_symbolic_trace.py", line 739, in trace
    (self.create_arg(fn(*args)),),
  File "/scratch/sanketpurandare/work/pytorch/torch/fx/_symbolic_trace.py", line 614, in flatten_fn
    tree_out = root_fn(*tree_args)
  File "/scratch/sanketpurandare/work/pytorch/torch/fx/experimental/proxy_tensor.py", line 422, in wrapped
    out = f(*tensors)
  File "/scratch/sanketpurandare/work/pytorch/functorch/_src/aot_autograd.py", line 166, in joint_forward_backward
    backward_out = torch.autograd.grad(
  File "/scratch/sanketpurandare/work/pytorch/torch/autograd/__init__.py", line 300, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/scratch/sanketpurandare/work/pytorch/torch/utils/_python_dispatch.py", line 101, in __torch_dispatch__
    return old.__torch_dispatch__(func, types, args, kwargs)
  File "/scratch/sanketpurandare/work/pytorch/torch/fx/experimental/proxy_tensor.py", line 448, in __torch_dispatch__
    return self.inner_torch_dispatch(func, types, args, kwargs)
  File "/scratch/sanketpurandare/work/pytorch/torch/fx/experimental/proxy_tensor.py", line 473, in inner_torch_dispatch
    out = proxy_call(self, func, args, kwargs)
  File "/scratch/sanketpurandare/work/pytorch/torch/fx/experimental/proxy_tensor.py", line 312, in proxy_call
    out = func(*args, **kwargs)
  File "/scratch/sanketpurandare/work/pytorch/torch/_ops.py", line 257, in __call__
    return self._op(*args, **kwargs or {})
  File "/scratch/sanketpurandare/work/pytorch/torch/utils/_python_dispatch.py", line 101, in __torch_dispatch__
    return old.__torch_dispatch__(func, types, args, kwargs)
  File "/scratch/sanketpurandare/work/pytorch/torch/_subclasses/fake_tensor.py", line 826, in __torch_dispatch__
    return run_fallback_kernel(
  File "/scratch/sanketpurandare/work/pytorch/torch/_subclasses/fake_tensor.py", line 949, in run_fallback_kernel
    r = func(*args, **kwargs)
  File "/scratch/sanketpurandare/work/pytorch/torch/_ops.py", line 257, in __call__
    return self._op(*args, **kwargs or {})
NotImplementedError: Could not run 'aten::select_backward' with arguments from the 'SparseCUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'aten::select_backward' is only available for these backends: [CPU, CUDA, HIP, MPS, IPU, XPU, HPU, VE, Meta, PrivateUse1, PrivateUse2, PrivateUse3, FPGA, ORT, Vulkan, Metal, QuantizedCPU, QuantizedCUDA, QuantizedHIP, QuantizedMPS, QuantizedIPU, QuantizedXPU, QuantizedHPU, QuantizedVE, QuantizedMeta, QuantizedPrivateUse1, QuantizedPrivateUse2, QuantizedPrivateUse3, CustomRNGKeyId, MkldnnCPU, SparseCsrCPU, SparseCsrCUDA, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradHIP, AutogradXLA, AutogradMPS, AutogradIPU, AutogradXPU, AutogradHPU, AutogradVE, AutogradLazy, AutogradMeta, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, AutogradNestedTensor, Tracer, AutocastCPU, AutocastCUDA, FuncTorchBatched, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PythonDispatcher].

Undefined: registered at /scratch/sanketpurandare/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel
```

### Versions

Collecting environment information...
PyTorch version: 1.14.0a0+git97de281
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.22.5
Libc version: glibc-2.27

Python version: 3.9.13 (main, Aug 25 2022, 23:26:10)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.0-1069-aws-x86_64-with-glibc2.27
Is CUDA available: True
CUDA runtime version: 11.6.112
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 510.47.03
cuDNN version: Probably one of the following:
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False

Versions of relevant libraries:
[pip3] bert-pytorch==0.0.1a4
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.1
[pip3] torch==1.14.0a0+git97de281
[pip3] torchaudio==0.12.0a0+5e96671
[pip3] torchmetrics==0.9.1
[pip3] torchrec-nightly==2022.4.26
[pip3] torchtext==0.14.0a0+8b35599
[pip3] torchvision==0.14.0a0+d1b2f4a
[pip3] torchx-nightly==2022.6.15
[conda] bert-pytorch              0.0.1a4                   dev_0    <develop>
[conda] blas                      1.0                         mkl  
[conda] magma-cuda116             2.6.1                         1    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-include               2022.0.1           h06a4308_117  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0  
[conda] numpy                     1.21.2                   pypi_0    pypi
[conda] numpy-base                1.23.1           py39ha15fc14_0  
[conda] torch                     1.14.0a0+git97de281           dev_0    <develop>
[conda] torchaudio                0.12.0a0+5e96671           dev_0    <develop>
[conda] torchdynamo               0.2.0                     dev_0    <develop>
[conda] torchmetrics              0.9.1                    pypi_0    pypi
[conda] torchrec-nightly          2022.4.26                pypi_0    pypi
[conda] torchtext                 0.14.0a0+8b35599           dev_0    <develop>
[conda] torchvision               0.14.0a0+d1b2f4a           dev_0    <develop>
[conda] torchx-nightly            2022.6.15                pypi_0    pyp

cc @zou3519 @Chillee @samdow @soumith