We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
`ProvenanceTensor` bug when used with `torch.set_default_device`
When trying to debug the code from the [forum post](https://forum.pyro.ai/t/running-pyro-plate-on-gpu/5097/8) I came across this bug:

```py
import torch
from pyro.ops.provenance import ProvenanceTensor

device = torch.device("cuda")
torch.set_default_device(device)
x = torch.tensor([1., 2., 3.])
y = ProvenanceTensor(x, frozenset(["x"]))
print(torch.as_tensor(y))
```
returns `tensor([], device='cuda:0')`

**Two observations:**
1. This doesn't happen if `torch.set_default_device` is not used: `print(torch.as_tensor(y.cuda()))` works fine
2. When `torch.set_default_device` is used then the following code gets invoked
https://github.com/pytorch/pytorch/blob/main/torch/utils/_device.py#L72-L76
where `func(*args, **kwargs)` returns `tensor([], device='cuda:0')`. This doesn't happen when `.cuda()` is used like in 1).

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pyro/ops/provenance.py`
Content:
```
1 # Copyright Contributors to the Pyro project.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 from functools import singledispatch
5 from typing import Tuple
6 
7 import torch
8 
9 
10 class ProvenanceTensor(torch.Tensor):
11     """
12     Provenance tracking implementation in Pytorch.
13 
14     This class wraps a :class:`torch.Tensor` to track provenance through
15     PyTorch ops, where provenance is a user-defined frozenset of objects. The
16     provenance of the output tensors of any op is the union of provenances of
17     input tensors.
18 
19     -   To start tracking provenance, wrap a :class:`torch.Tensor` in a
20         :class:`ProvenanceTensor` with user-defined initial provenance.
21     -   To read the provenance of a tensor use :meth:`get_provenance` .
22     -   To detach provenance during a computation (similar to
23         :meth:`~torch.Tensor.detach` to detach gradients during Pytorch
24         computations), use the :meth:`detach_provenance` . This is useful to
25         distinguish direct vs indirect provenance.
26 
27     Example::
28 
29         >>> a = ProvenanceTensor(torch.randn(3), frozenset({"a"}))
30         >>> b = ProvenanceTensor(torch.randn(3), frozenset({"b"}))
31         >>> c = torch.randn(3)
32         >>> assert get_provenance(a + b + c) == frozenset({"a", "b"})
33         >>> assert get_provenance(a + detach_provenance(b) + c) == frozenset({"a"})
34 
35     **References**
36 
37     [1] David Wingate, Noah Goodman, Andreas Stuhlmüller, Jeffrey Siskind (2011)
38         Nonstandard Interpretations of Probabilistic Programs for Efficient Inference
39         http://papers.neurips.cc/paper/4309-nonstandard-interpretations-of-probabilistic-programs-for-efficient-inference.pdf
40 
41     :param torch.Tensor data: An initial tensor to start tracking.
42     :param frozenset provenance: An initial provenance set.
43     """
44 
45     def __new__(cls, data: torch.Tensor, provenance=frozenset(), **kwargs):
46         assert not isinstance(data, ProvenanceTensor)
47         if not provenance:
48             return data
49         return super().__new__(cls)
50 
51     def __init__(self, data, provenance=frozenset()):
52         assert isinstance(provenance, frozenset)
53         if isinstance(data, ProvenanceTensor):
54             provenance |= data._provenance
55             data = data._t
56         self._t = data
57         self._provenance = provenance
58 
59     def __repr__(self):
60         return "Provenance:\n{}\nTensor:\n{}".format(self._provenance, self._t)
61 
62     @classmethod
63     def __torch_function__(cls, func, types, args=(), kwargs=None):
64         if kwargs is None:
65             kwargs = {}
66         # collect provenance information from args
67         provenance = frozenset()
68         # extract ProvenanceTensor._t data from args and kwargs
69         _args = []
70         for arg in args:
71             _arg, _provenance = extract_provenance(arg)
72             _args.append(_arg)
73             provenance |= _provenance
74         _kwargs = {}
75         for k, v in kwargs.items():
76             _v, _provenance = extract_provenance(v)
77             _kwargs[k] = _v
78             provenance |= provenance
79         ret = func(*_args, **_kwargs)
80         _ret = track_provenance(ret, provenance)
81         return _ret
82 
83 
84 @singledispatch
85 def track_provenance(x, provenance: frozenset):
86     """
87     Adds provenance info to the :class:`torch.Tensor` leaves of a data structure.
88 
89     :param x: an object to add provenence info to.
90     :param frozenset provenance: A provenence set.
91     :returns: A provenence-tracking version of ``x``.
92     """
93     return x
94 
95 
96 track_provenance.register(torch.Tensor)(ProvenanceTensor)
97 
98 
99 @track_provenance.register(frozenset)
100 @track_provenance.register(list)
101 @track_provenance.register(set)
102 @track_provenance.register(tuple)
103 def _track_provenance_list(x, provenance: frozenset):
104     return type(x)(track_provenance(part, provenance) for part in x)
105 
106 
107 @track_provenance.register
108 def _track_provenance_provenancetensor(x: ProvenanceTensor, provenance: frozenset):
109     x_value, old_provenance = extract_provenance(x)
110     return track_provenance(x_value, old_provenance | provenance)
111 
112 
113 @singledispatch
114 def extract_provenance(x) -> Tuple[object, frozenset]:
115     """
116     Extracts the provenance of a data structure possibly containing
117     :class:`torch.Tensor` s as leaves, and separates into a detached object and
118     provenance.
119 
120     :param x: An input data structure.
121     :returns: a tuple ``(detached_value, provenance)``
122     :rtype: tuple
123     """
124     return x, frozenset()
125 
126 
127 @extract_provenance.register(ProvenanceTensor)
128 def _extract_provenance_tensor(x):
129     return x._t, x._provenance
130 
131 
132 @extract_provenance.register(frozenset)
133 @extract_provenance.register(list)
134 @extract_provenance.register(set)
135 @extract_provenance.register(tuple)
136 def _extract_provenance_list(x):
137     provenance = frozenset()
138     values = []
139     for part in x:
140         v, p = extract_provenance(part)
141         values.append(v)
142         provenance |= p
143     value = type(x)(values)
144     return value, provenance
145 
146 
147 def get_provenance(x) -> frozenset:
148     """
149     Reads the provenance of a recursive datastructure possibly containing
150     :class:`torch.Tensor` s.
151 
152     :param torch.Tensor tensor: An input tensor.
153     :returns: A provenance frozenset.
154     :rtype: frozenset
155     """
156     _, provenance = extract_provenance(x)
157     return provenance
158 
159 
160 def detach_provenance(x):
161     """
162     Blocks provenance tracking through a tensor, similar to :meth:`torch.Tensor.detach`.
163 
164     :param torch.Tensor tensor: An input tensor.
165     :returns: A tensor sharing the same data but with no provenance.
166     :rtype: torch.Tensor
167     """
168     value, _ = extract_provenance(x)
169     return value
170
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pyro/ops/provenance.py b/pyro/ops/provenance.py
--- a/pyro/ops/provenance.py
+++ b/pyro/ops/provenance.py
@@ -46,14 +46,15 @@
         assert not isinstance(data, ProvenanceTensor)
         if not provenance:
             return data
-        return super().__new__(cls)
+        ret = data.as_subclass(cls)
+        ret._t = data  # this makes sure that detach_provenance always
+        # returns the same object. This is important when
+        # using the tensor as key in a dict, e.g. the global
+        # param store
+        return ret
 
     def __init__(self, data, provenance=frozenset()):
         assert isinstance(provenance, frozenset)
-        if isinstance(data, ProvenanceTensor):
-            provenance |= data._provenance
-            data = data._t
-        self._t = data
         self._provenance = provenance
 
     def __repr__(self):

{"golden_diff": "diff --git a/pyro/ops/provenance.py b/pyro/ops/provenance.py\n--- a/pyro/ops/provenance.py\n+++ b/pyro/ops/provenance.py\n@@ -46,14 +46,15 @@\n         assert not isinstance(data, ProvenanceTensor)\n         if not provenance:\n             return data\n-        return super().__new__(cls)\n+        ret = data.as_subclass(cls)\n+        ret._t = data  # this makes sure that detach_provenance always\n+        # returns the same object. This is important when\n+        # using the tensor as key in a dict, e.g. the global\n+        # param store\n+        return ret\n \n     def __init__(self, data, provenance=frozenset()):\n         assert isinstance(provenance, frozenset)\n-        if isinstance(data, ProvenanceTensor):\n-            provenance |= data._provenance\n-            data = data._t\n-        self._t = data\n         self._provenance = provenance\n \n     def __repr__(self):\n", "issue": "`ProvenanceTensor` bug when used with `torch.set_default_device`\nWhen trying to debug the code from the [forum post](https://forum.pyro.ai/t/running-pyro-plate-on-gpu/5097/8) I came across this bug:\r\n\r\n```py\r\nimport torch\r\nfrom pyro.ops.provenance import ProvenanceTensor\r\n\r\ndevice = torch.device(\"cuda\")\r\ntorch.set_default_device(device)\r\nx = torch.tensor([1., 2., 3.])\r\ny = ProvenanceTensor(x, frozenset([\"x\"]))\r\nprint(torch.as_tensor(y))\r\n```\r\nreturns `tensor([], device='cuda:0')`\r\n\r\n**Two observations:**\r\n1. This doesn't happen if `torch.set_default_device` is not used: `print(torch.as_tensor(y.cuda()))` works fine\r\n2. When `torch.set_default_device` is used then the following code gets invoked\r\nhttps://github.com/pytorch/pytorch/blob/main/torch/utils/_device.py#L72-L76\r\nwhere `func(*args, **kwargs)` returns `tensor([], device='cuda:0')`. This doesn't happen when `.cuda()` is used like in 1).\n", "before_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nfrom functools import singledispatch\nfrom typing import Tuple\n\nimport torch\n\n\nclass ProvenanceTensor(torch.Tensor):\n    \"\"\"\n    Provenance tracking implementation in Pytorch.\n\n    This class wraps a :class:`torch.Tensor` to track provenance through\n    PyTorch ops, where provenance is a user-defined frozenset of objects. The\n    provenance of the output tensors of any op is the union of provenances of\n    input tensors.\n\n    -   To start tracking provenance, wrap a :class:`torch.Tensor` in a\n        :class:`ProvenanceTensor` with user-defined initial provenance.\n    -   To read the provenance of a tensor use :meth:`get_provenance` .\n    -   To detach provenance during a computation (similar to\n        :meth:`~torch.Tensor.detach` to detach gradients during Pytorch\n        computations), use the :meth:`detach_provenance` . This is useful to\n        distinguish direct vs indirect provenance.\n\n    Example::\n\n        >>> a = ProvenanceTensor(torch.randn(3), frozenset({\"a\"}))\n        >>> b = ProvenanceTensor(torch.randn(3), frozenset({\"b\"}))\n        >>> c = torch.randn(3)\n        >>> assert get_provenance(a + b + c) == frozenset({\"a\", \"b\"})\n        >>> assert get_provenance(a + detach_provenance(b) + c) == frozenset({\"a\"})\n\n    **References**\n\n    [1] David Wingate, Noah Goodman, Andreas Stuhlm\u00fcller, Jeffrey Siskind (2011)\n        Nonstandard Interpretations of Probabilistic Programs for Efficient Inference\n        http://papers.neurips.cc/paper/4309-nonstandard-interpretations-of-probabilistic-programs-for-efficient-inference.pdf\n\n    :param torch.Tensor data: An initial tensor to start tracking.\n    :param frozenset provenance: An initial provenance set.\n    \"\"\"\n\n    def __new__(cls, data: torch.Tensor, provenance=frozenset(), **kwargs):\n        assert not isinstance(data, ProvenanceTensor)\n        if not provenance:\n            return data\n        return super().__new__(cls)\n\n    def __init__(self, data, provenance=frozenset()):\n        assert isinstance(provenance, frozenset)\n        if isinstance(data, ProvenanceTensor):\n            provenance |= data._provenance\n            data = data._t\n        self._t = data\n        self._provenance = provenance\n\n    def __repr__(self):\n        return \"Provenance:\\n{}\\nTensor:\\n{}\".format(self._provenance, self._t)\n\n    @classmethod\n    def __torch_function__(cls, func, types, args=(), kwargs=None):\n        if kwargs is None:\n            kwargs = {}\n        # collect provenance information from args\n        provenance = frozenset()\n        # extract ProvenanceTensor._t data from args and kwargs\n        _args = []\n        for arg in args:\n            _arg, _provenance = extract_provenance(arg)\n            _args.append(_arg)\n            provenance |= _provenance\n        _kwargs = {}\n        for k, v in kwargs.items():\n            _v, _provenance = extract_provenance(v)\n            _kwargs[k] = _v\n            provenance |= provenance\n        ret = func(*_args, **_kwargs)\n        _ret = track_provenance(ret, provenance)\n        return _ret\n\n\n@singledispatch\ndef track_provenance(x, provenance: frozenset):\n    \"\"\"\n    Adds provenance info to the :class:`torch.Tensor` leaves of a data structure.\n\n    :param x: an object to add provenence info to.\n    :param frozenset provenance: A provenence set.\n    :returns: A provenence-tracking version of ``x``.\n    \"\"\"\n    return x\n\n\ntrack_provenance.register(torch.Tensor)(ProvenanceTensor)\n\n\n@track_provenance.register(frozenset)\n@track_provenance.register(list)\n@track_provenance.register(set)\n@track_provenance.register(tuple)\ndef _track_provenance_list(x, provenance: frozenset):\n    return type(x)(track_provenance(part, provenance) for part in x)\n\n\n@track_provenance.register\ndef _track_provenance_provenancetensor(x: ProvenanceTensor, provenance: frozenset):\n    x_value, old_provenance = extract_provenance(x)\n    return track_provenance(x_value, old_provenance | provenance)\n\n\n@singledispatch\ndef extract_provenance(x) -> Tuple[object, frozenset]:\n    \"\"\"\n    Extracts the provenance of a data structure possibly containing\n    :class:`torch.Tensor` s as leaves, and separates into a detached object and\n    provenance.\n\n    :param x: An input data structure.\n    :returns: a tuple ``(detached_value, provenance)``\n    :rtype: tuple\n    \"\"\"\n    return x, frozenset()\n\n\n@extract_provenance.register(ProvenanceTensor)\ndef _extract_provenance_tensor(x):\n    return x._t, x._provenance\n\n\n@extract_provenance.register(frozenset)\n@extract_provenance.register(list)\n@extract_provenance.register(set)\n@extract_provenance.register(tuple)\ndef _extract_provenance_list(x):\n    provenance = frozenset()\n    values = []\n    for part in x:\n        v, p = extract_provenance(part)\n        values.append(v)\n        provenance |= p\n    value = type(x)(values)\n    return value, provenance\n\n\ndef get_provenance(x) -> frozenset:\n    \"\"\"\n    Reads the provenance of a recursive datastructure possibly containing\n    :class:`torch.Tensor` s.\n\n    :param torch.Tensor tensor: An input tensor.\n    :returns: A provenance frozenset.\n    :rtype: frozenset\n    \"\"\"\n    _, provenance = extract_provenance(x)\n    return provenance\n\n\ndef detach_provenance(x):\n    \"\"\"\n    Blocks provenance tracking through a tensor, similar to :meth:`torch.Tensor.detach`.\n\n    :param torch.Tensor tensor: An input tensor.\n    :returns: A tensor sharing the same data but with no provenance.\n    :rtype: torch.Tensor\n    \"\"\"\n    value, _ = extract_provenance(x)\n    return value\n", "path": "pyro/ops/provenance.py"}], "after_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nfrom functools import singledispatch\nfrom typing import Tuple\n\nimport torch\n\n\nclass ProvenanceTensor(torch.Tensor):\n    \"\"\"\n    Provenance tracking implementation in Pytorch.\n\n    This class wraps a :class:`torch.Tensor` to track provenance through\n    PyTorch ops, where provenance is a user-defined frozenset of objects. The\n    provenance of the output tensors of any op is the union of provenances of\n    input tensors.\n\n    -   To start tracking provenance, wrap a :class:`torch.Tensor` in a\n        :class:`ProvenanceTensor` with user-defined initial provenance.\n    -   To read the provenance of a tensor use :meth:`get_provenance` .\n    -   To detach provenance during a computation (similar to\n        :meth:`~torch.Tensor.detach` to detach gradients during Pytorch\n        computations), use the :meth:`detach_provenance` . This is useful to\n        distinguish direct vs indirect provenance.\n\n    Example::\n\n        >>> a = ProvenanceTensor(torch.randn(3), frozenset({\"a\"}))\n        >>> b = ProvenanceTensor(torch.randn(3), frozenset({\"b\"}))\n        >>> c = torch.randn(3)\n        >>> assert get_provenance(a + b + c) == frozenset({\"a\", \"b\"})\n        >>> assert get_provenance(a + detach_provenance(b) + c) == frozenset({\"a\"})\n\n    **References**\n\n    [1] David Wingate, Noah Goodman, Andreas Stuhlm\u00fcller, Jeffrey Siskind (2011)\n        Nonstandard Interpretations of Probabilistic Programs for Efficient Inference\n        http://papers.neurips.cc/paper/4309-nonstandard-interpretations-of-probabilistic-programs-for-efficient-inference.pdf\n\n    :param torch.Tensor data: An initial tensor to start tracking.\n    :param frozenset provenance: An initial provenance set.\n    \"\"\"\n\n    def __new__(cls, data: torch.Tensor, provenance=frozenset(), **kwargs):\n        assert not isinstance(data, ProvenanceTensor)\n        if not provenance:\n            return data\n        ret = data.as_subclass(cls)\n        ret._t = data  # this makes sure that detach_provenance always\n        # returns the same object. This is important when\n        # using the tensor as key in a dict, e.g. the global\n        # param store\n        return ret\n\n    def __init__(self, data, provenance=frozenset()):\n        assert isinstance(provenance, frozenset)\n        self._provenance = provenance\n\n    def __repr__(self):\n        return \"Provenance:\\n{}\\nTensor:\\n{}\".format(self._provenance, self._t)\n\n    @classmethod\n    def __torch_function__(cls, func, types, args=(), kwargs=None):\n        if kwargs is None:\n            kwargs = {}\n        # collect provenance information from args\n        provenance = frozenset()\n        # extract ProvenanceTensor._t data from args and kwargs\n        _args = []\n        for arg in args:\n            _arg, _provenance = extract_provenance(arg)\n            _args.append(_arg)\n            provenance |= _provenance\n        _kwargs = {}\n        for k, v in kwargs.items():\n            _v, _provenance = extract_provenance(v)\n            _kwargs[k] = _v\n            provenance |= provenance\n        ret = func(*_args, **_kwargs)\n        _ret = track_provenance(ret, provenance)\n        return _ret\n\n\n@singledispatch\ndef track_provenance(x, provenance: frozenset):\n    \"\"\"\n    Adds provenance info to the :class:`torch.Tensor` leaves of a data structure.\n\n    :param x: an object to add provenence info to.\n    :param frozenset provenance: A provenence set.\n    :returns: A provenence-tracking version of ``x``.\n    \"\"\"\n    return x\n\n\ntrack_provenance.register(torch.Tensor)(ProvenanceTensor)\n\n\n@track_provenance.register(frozenset)\n@track_provenance.register(list)\n@track_provenance.register(set)\n@track_provenance.register(tuple)\ndef _track_provenance_list(x, provenance: frozenset):\n    return type(x)(track_provenance(part, provenance) for part in x)\n\n\n@track_provenance.register\ndef _track_provenance_provenancetensor(x: ProvenanceTensor, provenance: frozenset):\n    x_value, old_provenance = extract_provenance(x)\n    return track_provenance(x_value, old_provenance | provenance)\n\n\n@singledispatch\ndef extract_provenance(x) -> Tuple[object, frozenset]:\n    \"\"\"\n    Extracts the provenance of a data structure possibly containing\n    :class:`torch.Tensor` s as leaves, and separates into a detached object and\n    provenance.\n\n    :param x: An input data structure.\n    :returns: a tuple ``(detached_value, provenance)``\n    :rtype: tuple\n    \"\"\"\n    return x, frozenset()\n\n\n@extract_provenance.register(ProvenanceTensor)\ndef _extract_provenance_tensor(x):\n    return x._t, x._provenance\n\n\n@extract_provenance.register(frozenset)\n@extract_provenance.register(list)\n@extract_provenance.register(set)\n@extract_provenance.register(tuple)\ndef _extract_provenance_list(x):\n    provenance = frozenset()\n    values = []\n    for part in x:\n        v, p = extract_provenance(part)\n        values.append(v)\n        provenance |= p\n    value = type(x)(values)\n    return value, provenance\n\n\ndef get_provenance(x) -> frozenset:\n    \"\"\"\n    Reads the provenance of a recursive datastructure possibly containing\n    :class:`torch.Tensor` s.\n\n    :param torch.Tensor tensor: An input tensor.\n    :returns: A provenance frozenset.\n    :rtype: frozenset\n    \"\"\"\n    _, provenance = extract_provenance(x)\n    return provenance\n\n\ndef detach_provenance(x):\n    \"\"\"\n    Blocks provenance tracking through a tensor, similar to :meth:`torch.Tensor.detach`.\n\n    :param torch.Tensor tensor: An input tensor.\n    :returns: A tensor sharing the same data but with no provenance.\n    :rtype: torch.Tensor\n    \"\"\"\n    value, _ = extract_provenance(x)\n    return value\n", "path": "pyro/ops/provenance.py"}]}