prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Memory leak using TraceEnum_ELBO
  I noticed a major memory leak when training SVI using `TraceEnum_ELBO`.
  I initially noticed this in a custom model we are developing but then I found it seems a more general bug.
  
  For example, it affects even the Pyro tutorials GMM example [here](https://pyro.ai/examples/gmm.html). Where memory usage rapidly goes from a couple of hundred MBs to a many **GB**s very quickly!
  
  I have run this Macbook Pro 2019 running MacOS 10.15. To replicate the issue is enough running the notebook linked.
  
  I have tried to comment out the following lines and add a garbage collector call, that reduces the entity of the memory accumulation of one order of magnitude but does not solve the problem completely, which becomes particularly severe for large datasets.
  
  ```python
  # Register hooks to monitor gradient norms.
  # gradient_norms = defaultdict(list)
  # for name, value in pyro.get_param_store().named_parameters():
  #     value.register_hook(lambda g, name=name: gradient_norms[name].append(g.norm().item()))
  
  import gc
  losses = []
  for i in range(200000):
      loss = svi.step(data)
      #losses.append(loss)
      gc.collect()
  ``` 
  
  (from [this forum post](https://forum.pyro.ai/t/memory-usage-grows-with-iterations-in-svi-step/4103))
  [bug] Memory leak on GPU
  ### Issue Description
  After successfully following the SS-VAE example in pyro I have moved over to training a variant of it on a GPU. A memory leak has arisen that has been boiled down to a minimum working example that replicates the error. I am at a loss as to how to proceed in rectifying this issue, any help or insight would be much appreciated!
  
  ### Environment
   - OS and python version. 
     - OS: Red Hat Enterprise Linux Workstation release 7.9 (Maipo) 
     - Python: 3.9.7
     - CUDA Version: 11.2
   - PyTorch version: 1.10.0
   - Pyro version: 1.7.0
  
  ### Code Snippet
  ``` python
  from datetime import datetime
  
  import pyro
  import pyro.distributions as dist
  import torch
  import torch.nn as nn
  from pyro.infer import SVI, TraceEnum_ELBO
  from pyro.optim import Adam
  
  
  class SSVAE(nn.Module):
      def __init__(self):
          super().__init__()
          self.accept_ratio = .5
          ones = torch.ones((), device='cuda')
          ratio = 1 / (self.accept_ratio + 1)
          self.alpha_prior = torch.stack((ones - ratio, ones - self.accept_ratio * ratio)).T
          self.cuda()
  
      def model(self):
          pyro.sample("y", dist.OneHotCategorical(self.alpha_prior))
  
      def guide(self):
          pass
  
  
  def sizeof_fmt(num, suffix="B"):
      for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]:
          if abs(num) < 1024.0:
              return f"{num:3.1f}{unit}{suffix}"
          num /= 1024.0
      return f"{num:.1f}Yi{suffix}"
  
  
  if __name__ == '__main__':
      torch.cuda.empty_cache()
      ssvae = SSVAE()
  
      optimizer = Adam({"lr": 1e-4})
  
      svi = SVI(ssvae.model, ssvae.guide, optimizer, loss=TraceEnum_ELBO())
  
      for i in range(10):
          print(f"{datetime.now().strftime('%H:%M:%S')} - Starting epoch #{i + 1}")
          print(
              f"Current GPU memory allocation:"
              f" {sizeof_fmt(torch.cuda.memory_allocated())}/{sizeof_fmt(torch.cuda.max_memory_allocated())}")
  
          # get the losses for an epoch
          for _ in range(100):
              svi.step()
  ```
  
  
  --- 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/einsum/torch_log.py`
  Content:
  ```
  1 # Copyright (c) 2017-2019 Uber Technologies, Inc.
  2 # SPDX-License-Identifier: Apache-2.0
  3 
  4 import torch
  5 
  6 from pyro.ops.einsum.util import Tensordot
  7 from pyro.ops.special import safe_log
  8 
  9 
  10 def transpose(a, axes):
  11     return a.permute(axes)
  12 
  13 
  14 def einsum(equation, *operands):
  15     """
  16     Log-sum-exp implementation of einsum.
  17     """
  18     # rename symbols to support PyTorch 0.4.1 and earlier,
  19     # which allow only symbols a-z.
  20     symbols = sorted(set(equation) - set(",->"))
  21     rename = dict(zip(symbols, "abcdefghijklmnopqrstuvwxyz"))
  22     equation = "".join(rename.get(s, s) for s in equation)
  23 
  24     inputs, output = equation.split("->")
  25     if inputs == output:
  26         return operands[0][...]  # create a new object
  27     inputs = inputs.split(",")
  28 
  29     shifts = []
  30     exp_operands = []
  31     for dims, operand in zip(inputs, operands):
  32         shift = operand.detach()
  33         for i, dim in enumerate(dims):
  34             if dim not in output:
  35                 shift = shift.max(i, keepdim=True)[0]
  36         # avoid nan due to -inf - -inf
  37         shift = shift.clamp(min=torch.finfo(shift.dtype).min)
  38         exp_operands.append((operand - shift).exp())
  39 
  40         # permute shift to match output
  41         shift = shift.reshape(
  42             torch.Size(size for size, dim in zip(operand.shape, dims) if dim in output)
  43         )
  44         if shift.dim():
  45             shift = shift.reshape((1,) * (len(output) - shift.dim()) + shift.shape)
  46             dims = [dim for dim in dims if dim in output]
  47             dims = [dim for dim in output if dim not in dims] + dims
  48             shift = shift.permute(*(dims.index(dim) for dim in output))
  49         shifts.append(shift)
  50 
  51     result = safe_log(torch.einsum(equation, exp_operands))
  52     return sum(shifts + [result])
  53 
  54 
  55 tensordot = Tensordot(einsum)
  56 
  57 __all__ = ["transpose", "einsum", "tensordot"]
  58
  ```
  --- 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()
  
  ```
  
golden_diff: diff --git a/pyro/ops/einsum/torch_log.py b/pyro/ops/einsum/torch_log.py
  --- a/pyro/ops/einsum/torch_log.py
  +++ b/pyro/ops/einsum/torch_log.py
  @@ -23,7 +23,10 @@
   
       inputs, output = equation.split("->")
       if inputs == output:
  -        return operands[0][...]  # create a new object
  +        # Originally we return `operands[0][...]` but that caused
  +        # memory leak in PyTorch >= 1.11 (issue #3068). Hence we
  +        # return `operands[0].clone()` here.
  +        return operands[0].clone()  # create a new object
       inputs = inputs.split(",")
   
       shifts = []
  
verification_info: {"golden_diff": "diff --git a/pyro/ops/einsum/torch_log.py b/pyro/ops/einsum/torch_log.py\n--- a/pyro/ops/einsum/torch_log.py\n+++ b/pyro/ops/einsum/torch_log.py\n@@ -23,7 +23,10 @@\n \n     inputs, output = equation.split(\"->\")\n     if inputs == output:\n-        return operands[0][...]  # create a new object\n+        # Originally we return `operands[0][...]` but that caused\n+        # memory leak in PyTorch >= 1.11 (issue #3068). Hence we\n+        # return `operands[0].clone()` here.\n+        return operands[0].clone()  # create a new object\n     inputs = inputs.split(\",\")\n \n     shifts = []\n", "issue": "Memory leak using TraceEnum_ELBO\nI noticed a major memory leak when training SVI using `TraceEnum_ELBO`.\r\nI initially noticed this in a custom model we are developing but then I found it seems a more general bug.\r\n\r\nFor example, it affects even the Pyro tutorials GMM example [here](https://pyro.ai/examples/gmm.html). Where memory usage rapidly goes from a couple of hundred MBs to a many **GB**s very quickly!\r\n\r\nI have run this Macbook Pro 2019 running MacOS 10.15. To replicate the issue is enough running the notebook linked.\r\n\r\nI have tried to comment out the following lines and add a garbage collector call, that reduces the entity of the memory accumulation of one order of magnitude but does not solve the problem completely, which becomes particularly severe for large datasets.\r\n\r\n```python\r\n# Register hooks to monitor gradient norms.\r\n# gradient_norms = defaultdict(list)\r\n# for name, value in pyro.get_param_store().named_parameters():\r\n#     value.register_hook(lambda g, name=name: gradient_norms[name].append(g.norm().item()))\r\n\r\nimport gc\r\nlosses = []\r\nfor i in range(200000):\r\n    loss = svi.step(data)\r\n    #losses.append(loss)\r\n    gc.collect()\r\n``` \r\n\r\n(from [this forum post](https://forum.pyro.ai/t/memory-usage-grows-with-iterations-in-svi-step/4103))\n[bug] Memory leak on GPU\n### Issue Description\r\nAfter successfully following the SS-VAE example in pyro I have moved over to training a variant of it on a GPU. A memory leak has arisen that has been boiled down to a minimum working example that replicates the error. I am at a loss as to how to proceed in rectifying this issue, any help or insight would be much appreciated!\r\n\r\n### Environment\r\n - OS and python version. \r\n   - OS: Red Hat Enterprise Linux Workstation release 7.9 (Maipo) \r\n   - Python: 3.9.7\r\n   - CUDA Version: 11.2\r\n - PyTorch version: 1.10.0\r\n - Pyro version: 1.7.0\r\n\r\n### Code Snippet\r\n``` python\r\nfrom datetime import datetime\r\n\r\nimport pyro\r\nimport pyro.distributions as dist\r\nimport torch\r\nimport torch.nn as nn\r\nfrom pyro.infer import SVI, TraceEnum_ELBO\r\nfrom pyro.optim import Adam\r\n\r\n\r\nclass SSVAE(nn.Module):\r\n    def __init__(self):\r\n        super().__init__()\r\n        self.accept_ratio = .5\r\n        ones = torch.ones((), device='cuda')\r\n        ratio = 1 / (self.accept_ratio + 1)\r\n        self.alpha_prior = torch.stack((ones - ratio, ones - self.accept_ratio * ratio)).T\r\n        self.cuda()\r\n\r\n    def model(self):\r\n        pyro.sample(\"y\", dist.OneHotCategorical(self.alpha_prior))\r\n\r\n    def guide(self):\r\n        pass\r\n\r\n\r\ndef sizeof_fmt(num, suffix=\"B\"):\r\n    for unit in [\"\", \"Ki\", \"Mi\", \"Gi\", \"Ti\", \"Pi\", \"Ei\", \"Zi\"]:\r\n        if abs(num) < 1024.0:\r\n            return f\"{num:3.1f}{unit}{suffix}\"\r\n        num /= 1024.0\r\n    return f\"{num:.1f}Yi{suffix}\"\r\n\r\n\r\nif __name__ == '__main__':\r\n    torch.cuda.empty_cache()\r\n    ssvae = SSVAE()\r\n\r\n    optimizer = Adam({\"lr\": 1e-4})\r\n\r\n    svi = SVI(ssvae.model, ssvae.guide, optimizer, loss=TraceEnum_ELBO())\r\n\r\n    for i in range(10):\r\n        print(f\"{datetime.now().strftime('%H:%M:%S')} - Starting epoch #{i + 1}\")\r\n        print(\r\n            f\"Current GPU memory allocation:\"\r\n            f\" {sizeof_fmt(torch.cuda.memory_allocated())}/{sizeof_fmt(torch.cuda.max_memory_allocated())}\")\r\n\r\n        # get the losses for an epoch\r\n        for _ in range(100):\r\n            svi.step()\r\n```\r\n\n", "before_files": [{"content": "# Copyright (c) 2017-2019 Uber Technologies, Inc.\n# SPDX-License-Identifier: Apache-2.0\n\nimport torch\n\nfrom pyro.ops.einsum.util import Tensordot\nfrom pyro.ops.special import safe_log\n\n\ndef transpose(a, axes):\n    return a.permute(axes)\n\n\ndef einsum(equation, *operands):\n    \"\"\"\n    Log-sum-exp implementation of einsum.\n    \"\"\"\n    # rename symbols to support PyTorch 0.4.1 and earlier,\n    # which allow only symbols a-z.\n    symbols = sorted(set(equation) - set(\",->\"))\n    rename = dict(zip(symbols, \"abcdefghijklmnopqrstuvwxyz\"))\n    equation = \"\".join(rename.get(s, s) for s in equation)\n\n    inputs, output = equation.split(\"->\")\n    if inputs == output:\n        return operands[0][...]  # create a new object\n    inputs = inputs.split(\",\")\n\n    shifts = []\n    exp_operands = []\n    for dims, operand in zip(inputs, operands):\n        shift = operand.detach()\n        for i, dim in enumerate(dims):\n            if dim not in output:\n                shift = shift.max(i, keepdim=True)[0]\n        # avoid nan due to -inf - -inf\n        shift = shift.clamp(min=torch.finfo(shift.dtype).min)\n        exp_operands.append((operand - shift).exp())\n\n        # permute shift to match output\n        shift = shift.reshape(\n            torch.Size(size for size, dim in zip(operand.shape, dims) if dim in output)\n        )\n        if shift.dim():\n            shift = shift.reshape((1,) * (len(output) - shift.dim()) + shift.shape)\n            dims = [dim for dim in dims if dim in output]\n            dims = [dim for dim in output if dim not in dims] + dims\n            shift = shift.permute(*(dims.index(dim) for dim in output))\n        shifts.append(shift)\n\n    result = safe_log(torch.einsum(equation, exp_operands))\n    return sum(shifts + [result])\n\n\ntensordot = Tensordot(einsum)\n\n__all__ = [\"transpose\", \"einsum\", \"tensordot\"]\n", "path": "pyro/ops/einsum/torch_log.py"}], "after_files": [{"content": "# Copyright (c) 2017-2019 Uber Technologies, Inc.\n# SPDX-License-Identifier: Apache-2.0\n\nimport torch\n\nfrom pyro.ops.einsum.util import Tensordot\nfrom pyro.ops.special import safe_log\n\n\ndef transpose(a, axes):\n    return a.permute(axes)\n\n\ndef einsum(equation, *operands):\n    \"\"\"\n    Log-sum-exp implementation of einsum.\n    \"\"\"\n    # rename symbols to support PyTorch 0.4.1 and earlier,\n    # which allow only symbols a-z.\n    symbols = sorted(set(equation) - set(\",->\"))\n    rename = dict(zip(symbols, \"abcdefghijklmnopqrstuvwxyz\"))\n    equation = \"\".join(rename.get(s, s) for s in equation)\n\n    inputs, output = equation.split(\"->\")\n    if inputs == output:\n        # Originally we return `operands[0][...]` but that caused\n        # memory leak in PyTorch >= 1.11 (issue #3068). Hence we\n        # return `operands[0].clone()` here.\n        return operands[0].clone()  # create a new object\n    inputs = inputs.split(\",\")\n\n    shifts = []\n    exp_operands = []\n    for dims, operand in zip(inputs, operands):\n        shift = operand.detach()\n        for i, dim in enumerate(dims):\n            if dim not in output:\n                shift = shift.max(i, keepdim=True)[0]\n        # avoid nan due to -inf - -inf\n        shift = shift.clamp(min=torch.finfo(shift.dtype).min)\n        exp_operands.append((operand - shift).exp())\n\n        # permute shift to match output\n        shift = shift.reshape(\n            torch.Size(size for size, dim in zip(operand.shape, dims) if dim in output)\n        )\n        if shift.dim():\n            shift = shift.reshape((1,) * (len(output) - shift.dim()) + shift.shape)\n            dims = [dim for dim in dims if dim in output]\n            dims = [dim for dim in output if dim not in dims] + dims\n            shift = shift.permute(*(dims.index(dim) for dim in output))\n        shifts.append(shift)\n\n    result = safe_log(torch.einsum(equation, exp_operands))\n    return sum(shifts + [result])\n\n\ntensordot = Tensordot(einsum)\n\n__all__ = [\"transpose\", \"einsum\", \"tensordot\"]\n", "path": "pyro/ops/einsum/torch_log.py"}]}