Repository: pytorch
Issue: Unable to backprop through dense weighted sum of sparse_coo_tensors
Labels: module: sparse,triaged

### 🐛 Describe the bug

I can't find a way to perform a weighted sum of `sparse_coo_tensors` using a dense tensor parameter without running into an `aten::sum.dim_IntList` NotImplementedError. This seems like a common use-case, so I was wondering if there's a workaround, or correct way to accomplish this operation.

```python
import torch

sp_mtx1 = torch.eye(10, dtype=float).to_sparse_coo().cuda()
sp_mtx2 = torch.eye(10, dtype=float).to_sparse_coo().cuda()
sp_mtx_comb = torch.stack([sp_mtx1, sp_mtx2], 0) # (2, 10, 10)
weight = torch.rand([2], dtype=float, requires_grad=True).cuda()

weighted = sp_mtx_comb*weight.view(-1,1,1) # (2, 10, 10)
weighted_sum  = torch.sparse.sum(weighted, 0) # (10, 10)
loss = torch.sparse.sum(weighted_sum)
loss.backward()
```

This example produces the following traceback:
```
    def test_reproduce_cuda_config_err():
        sp_mtx0 = torch.eye(10, dtype=float).to_sparse_coo().cuda()
        sp_mtx1 = torch.eye(10, dtype=float).to_sparse_coo().cuda()
        sp_mtx = torch.stack([sp_mtx0, sp_mtx1], 0)
        weight = torch.tensor([1.0, 2.0], dtype=float, requires_grad=True).cuda()
    
        res = sp_mtx*weight.view(-1,1,1)
        print(res)
        weighted_sum = torch.sparse.sum(res, 0)
        print(weighted_sum)
        loss = torch.sparse.sum(weighted_sum)
>       loss.backward()

test_reproduce.py:17: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
../../.conda/envs/torch113/lib/python3.10/site-packages/torch/_tensor.py:487: in backward
    torch.autograd.backward(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

tensors = (tensor(30., device='cuda:0', dtype=torch.float64, grad_fn=<SumBackward0>),), grad_tensors = None, retain_graph = False, create_graph = False, grad_variables = None
inputs = ()

    def backward(
        tensors: _TensorOrTensors,
        grad_tensors: Optional[_TensorOrTensors] = None,
        retain_graph: Optional[bool] = None,
        create_graph: bool = False,
        grad_variables: Optional[_TensorOrTensors] = None,
        inputs: Optional[_TensorOrTensors] = None,
    ) -> None:
        r"""Computes the sum of gradients of given tensors with respect to graph
        leaves.
    
        The graph is differentiated using the chain rule. If any of ``tensors``
        are non-scalar (i.e. their data has more than one element) and require
        gradient, then the Jacobian-vector product would be computed, in this
        case the function additionally requires specifying ``grad_tensors``.
        It should be a sequence of matching length, that contains the "vector"
        in the Jacobian-vector product, usually the gradient of the differentiated
        function w.r.t. corresponding tensors (``None`` is an acceptable value for
        all tensors that don't need gradient tensors).
    
        This function accumulates gradients in the leaves - you might need to zero
        ``.grad`` attributes or set them to ``None`` before calling it.
        See :ref:`Default gradient layouts<default-grad-layouts>`
        for details on the memory layout of accumulated gradients.
    
        .. note::
            Using this method with ``create_graph=True`` will create a reference cycle
            between the parameter and its gradient which can cause a memory leak.
            We recommend using ``autograd.grad`` when creating the graph to avoid this.
            If you have to use this function, make sure to reset the ``.grad`` fields of your
            parameters to ``None`` after use to break the cycle and avoid the leak.
    
        .. note::
    
            If you run any forward ops, create ``grad_tensors``, and/or call ``backward``
            in a user-specified CUDA stream context, see
            :ref:`Stream semantics of backward passes<bwd-cuda-stream-semantics>`.
    
        .. note::
    
            When ``inputs`` are provided and a given input is not a leaf,
            the current implementation will call its grad_fn (even though it is not strictly needed to get this gradients).
            It is an implementation detail on which the user should not rely.
            See https://github.com/pytorch/pytorch/pull/60521#issuecomment-867061780 for more details.
    
        Args:
            tensors (Sequence[Tensor] or Tensor): Tensors of which the derivative will be
                computed.
            grad_tensors (Sequence[Tensor or None] or Tensor, optional): The "vector" in
                the Jacobian-vector product, usually gradients w.r.t. each element of
                corresponding tensors. None values can be specified for scalar Tensors or
                ones that don't require grad. If a None value would be acceptable for all
                grad_tensors, then this argument is optional.
            retain_graph (bool, optional): If ``False``, the graph used to compute the grad
                will be freed. Note that in nearly all cases setting this option to ``True``
                is not needed and often can be worked around in a much more efficient
                way. Defaults to the value of ``create_graph``.
            create_graph (bool, optional): If ``True``, graph of the derivative will
                be constructed, allowing to compute higher order derivative products.
                Defaults to ``False``.
            inputs (Sequence[Tensor] or Tensor, optional): Inputs w.r.t. which the gradient
                be will accumulated into ``.grad``. All other Tensors will be ignored. If
                not provided, the gradient is accumulated into all the leaf Tensors that
                were used to compute the attr::tensors.
        """
        if torch._C._are_functorch_transforms_active():
            raise RuntimeError(
                "backward() called inside a functorch transform. This is not "
                "supported, please use functorch.grad or functorch.vjp instead "
                "or call backward() outside of functorch transforms.")
    
        if grad_variables is not None:
            warnings.warn("'grad_variables' is deprecated. Use 'grad_tensors' instead.")
            if grad_tensors is None:
                grad_tensors = grad_variables
            else:
                raise RuntimeError("'grad_tensors' and 'grad_variables' (deprecated) "
                                   "arguments both passed to backward(). Please only "
                                   "use 'grad_tensors'.")
        if inputs is not None and len(inputs) == 0:
            raise RuntimeError("'inputs' argument to backward() cannot be empty.")
    
        tensors = (tensors,) if isinstance(tensors, torch.Tensor) else tuple(tensors)
        inputs = (inputs,) if isinstance(inputs, torch.Tensor) else \
            tuple(inputs) if inputs is not None else tuple()
    
        grad_tensors_ = _tensor_or_tensors_to_tuple(grad_tensors, len(tensors))
        grad_tensors_ = _make_grads(tensors, grad_tensors_, is_grads_batched=False)
        if retain_graph is None:
            retain_graph = create_graph
    
        # The reason we repeat same the comment below is that
        # some Python versions print out the first line of a multi-line function
        # calls in the traceback and some print out the last line
>       Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
            tensors, grad_tensors_, retain_graph, create_graph, inputs,
            allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to run the backward pass
E       NotImplementedError: Could not run 'aten::sum.dim_IntList' 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::sum.dim_IntList' 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, NestedTensorCPU, 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].
E       
E       Undefined: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       CPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCPU.cpp:30798 [kernel]
E       CUDA: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCUDA.cpp:43635 [kernel]
E       HIP: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       MPS: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       IPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       XPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       HPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       VE: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       Meta: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterMeta.cpp:26815 [kernel]
E       PrivateUse1: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       PrivateUse2: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       PrivateUse3: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       FPGA: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       ORT: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       Vulkan: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       Metal: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       QuantizedCPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       QuantizedCUDA: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       QuantizedHIP: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       QuantizedMPS: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       QuantizedIPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       QuantizedXPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       QuantizedHPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       QuantizedVE: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       QuantizedMeta: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       QuantizedPrivateUse1: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       QuantizedPrivateUse2: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       QuantizedPrivateUse3: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       CustomRNGKeyId: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       MkldnnCPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       SparseCsrCPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       SparseCsrCUDA: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp:21389 [default backend kernel]
E       NestedTensorCPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/build/aten/src/ATen/RegisterNestedTensorCPU.cpp:457 [kernel]
E       BackendSelect: fallthrough registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/core/BackendSelectFallbackKernel.cpp:3 [backend fallback]
E       Python: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/core/PythonFallbackKernel.cpp:140 [backend fallback]
E       FuncTorchDynamicLayerBackMode: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/functorch/DynamicLayer.cpp:488 [backend fallback]
E       Functionalize: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/FunctionalizeFallbackKernel.cpp:291 [backend fallback]
E       Named: fallthrough registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/core/NamedRegistrations.cpp:11 [kernel]
E       Conjugate: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/ConjugateFallback.cpp:18 [backend fallback]
E       Negative: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/native/NegateFallback.cpp:18 [backend fallback]
E       ZeroTensor: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/ZeroTensorFallback.cpp:86 [backend fallback]
E       ADInplaceOrView: fallthrough registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/core/VariableFallbackKernel.cpp:64 [backend fallback]
E       AutogradOther: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradCPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradCUDA: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradHIP: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradXLA: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradMPS: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradIPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradXPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradHPU: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradVE: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradLazy: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradMeta: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradPrivateUse1: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradPrivateUse2: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradPrivateUse3: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:16899 [autograd kernel]
E       AutogradNestedTensor: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/VariableType_2.cpp:17970 [kernel]
E       Tracer: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/torch/csrc/autograd/generated/TraceType_2.cpp:16890 [kernel]
E       AutocastCPU: fallthrough registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/autocast_mode.cpp:482 [backend fallback]
E       AutocastCUDA: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/autocast_mode.cpp:328 [kernel]
E       FuncTorchBatched: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/functorch/BatchRulesReduceOps.cpp:371 [kernel]
E       FuncTorchVmapMode: fallthrough registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/functorch/VmapModeRegistrations.cpp:28 [backend fallback]
E       Batched: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/BatchingRegistrations.cpp:1068 [kernel]
E       VmapMode: fallthrough registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/VmapModeRegistrations.cpp:33 [backend fallback]
E       FuncTorchGradWrapper: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/functorch/TensorWrapper.cpp:189 [backend fallback]
E       PythonTLSSnapshot: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/core/PythonFallbackKernel.cpp:148 [backend fallback]
E       FuncTorchDynamicLayerFrontMode: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/functorch/DynamicLayer.cpp:484 [backend fallback]
E       PythonDispatcher: registered at /opt/conda/conda-bld/pytorch_1666643003845/work/aten/src/ATen/core/PythonFallbackKernel.cpp:144 [backend fallback]

../../.conda/envs/torch113/lib/python3.10/site-packages/torch/autograd/__init__.py:197: NotImplementedError
```

### Versions

PyTorch version: 1.13.0
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: NAME="Red Hat Enterprise Linux Server" (x86_64)
GCC version: (GCC) 4.8.5 20150623 (Red Hat 4.8.5-44)
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.17

Python version: 3.10.8 (main, Nov  4 2022, 13:48:29) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-3.10.0-1160.59.1.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.6.124
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA TITAN RTX
Nvidia driver version: 510.47.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.3
[pip3] torch==1.13.0
[pip3] torchaudio==0.13.0
[pip3] torchvision==0.14.0
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-fft                   1.3.1                    pypi_0    pypi
[conda] mkl-random                1.2.2                    pypi_0    pypi
[conda] mkl-service               2.4.0                    pypi_0    pypi
[conda] mkl_fft                   1.3.1           py310hd6ae3a3_0  
[conda] mkl_random                1.2.2           py310h00e6091_0  
[conda] numpy                     1.23.3                   pypi_0    pypi
[conda] numpy-base                1.23.3          py310h8e6c178_1  
[conda] pytorch                   1.13.0          py3.10_cuda11.6_cudnn8.3.2_0    pytorch
[conda] pytorch-cuda              11.6                 h867d48c_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch                     1.13.0                   pypi_0    pypi
[conda] torchaudio                0.13.0                   pypi_0    pypi
[conda] torchvision               0.14.0                   pypi_0    pypi

cc @nikitaved @pearu @cpuhrsch @amjames @bhosmer