prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  JIT leaks memory in nets with dropout layer 
  if I add a dropout layer in FC layers, a JIT compiled forward pass of a net starts to seriously leak memory in the backward pass. I made a simple test case below. When it runs, you will see a rapidly increasing GPU/CPU memory usage in nvidia-smi/top and it will crash with an out of memory exception soon. Without the dropout layer or the backward pass, it works fine. 
  
  ```python
  import torch
  from torch import jit
  import torch.nn as nn
  from torch.autograd import Variable
  
  class TestNet(nn.Module):
      def __init__(self):
          super(TestNet, self).__init__()
          self.net1 = nn.Linear(100, 200)
          self.net2 = nn.Linear(200, 1)
          self.sigmoid = nn.Sigmoid()
          self.ReLU = nn.ReLU(inplace=False)
          self.drop = nn.Dropout(0.5)
                 
      def forward(self, V):
          return self.sigmoid(self.net2(self.drop(self.ReLU(self.net1(V))))).squeeze() 
  
  
  use_cuda = True
  net = TestNet()
  criterion = nn.BCELoss()
  if use_cuda:
      net.cuda()
      criterion.cuda()
      V = Variable(torch.randn(100, 100)).cuda()
      label = Variable(torch.randn(100)).cuda()
  else:
      V = Variable(torch.randn(100, 100))
      label = Variable(torch.randn(100))
  
  net.train()
  fwd = jit.compile(net.forward)
  for i in range(0,1000000):
      r = fwd(V)
      err = criterion(r, label)
      err.backward()   
  
  ```
  
  
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `torch/nn/_functions/dropout.py`
  Content:
  ```
  1 import torch
  2 from torch.autograd.function import InplaceFunction
  3 from torch.autograd import Variable
  4 from itertools import repeat
  5 
  6 
  7 class Dropout(InplaceFunction):
  8 
  9     @staticmethod
  10     def _make_noise(input):
  11         return input.new().resize_as_(input)
  12 
  13     @staticmethod
  14     def symbolic(g, input, p=0.5, train=False, inplace=False):
  15         # See Note [Export inplace]
  16         r, _ = g.op("Dropout", input, ratio_f=p, is_test_i=not train, outputs=2)
  17         return r
  18 
  19     @classmethod
  20     def forward(cls, ctx, input, p=0.5, train=False, inplace=False):
  21         if p < 0 or p > 1:
  22             raise ValueError("dropout probability has to be between 0 and 1, "
  23                              "but got {}".format(p))
  24         ctx.p = p
  25         ctx.train = train
  26         ctx.inplace = inplace
  27 
  28         if ctx.inplace:
  29             ctx.mark_dirty(input)
  30             output = input
  31         else:
  32             output = input.clone()
  33 
  34         if ctx.p > 0 and ctx.train:
  35             ctx.noise = cls._make_noise(input)
  36             if ctx.p == 1:
  37                 ctx.noise.fill_(0)
  38             else:
  39                 ctx.noise.bernoulli_(1 - ctx.p).div_(1 - ctx.p)
  40             ctx.noise = ctx.noise.expand_as(input)
  41             output.mul_(ctx.noise)
  42 
  43         return output
  44 
  45     @staticmethod
  46     def backward(ctx, grad_output):
  47         if ctx.p > 0 and ctx.train:
  48             return grad_output.mul(Variable(ctx.noise)), None, None, None
  49         else:
  50             return grad_output, None, None, None
  51 
  52 
  53 class FeatureDropout(Dropout):
  54 
  55     @staticmethod
  56     def symbolic(g, input, p=0.5, train=False, inplace=False):
  57         # See Note [Export inplace]
  58         # NB: In inference mode, FeatureDropout is exported as an identity op.
  59         from torch.onnx.symbolic import _unimplemented
  60         if train:
  61             return _unimplemented("FeatureDropout", "training mode")
  62         return input
  63 
  64     @staticmethod
  65     def _make_noise(input):
  66         return input.new().resize_(input.size(0), input.size(1),
  67                                    *repeat(1, input.dim() - 2))
  68
  ```
  --- 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/torch/nn/_functions/dropout.py b/torch/nn/_functions/dropout.py
  --- a/torch/nn/_functions/dropout.py
  +++ b/torch/nn/_functions/dropout.py
  @@ -25,20 +25,22 @@
           ctx.train = train
           ctx.inplace = inplace
   
  +        if ctx.p == 0 or not ctx.train:
  +            return input
  +
           if ctx.inplace:
               ctx.mark_dirty(input)
               output = input
           else:
               output = input.clone()
   
  -        if ctx.p > 0 and ctx.train:
  -            ctx.noise = cls._make_noise(input)
  -            if ctx.p == 1:
  -                ctx.noise.fill_(0)
  -            else:
  -                ctx.noise.bernoulli_(1 - ctx.p).div_(1 - ctx.p)
  -            ctx.noise = ctx.noise.expand_as(input)
  -            output.mul_(ctx.noise)
  +        ctx.noise = cls._make_noise(input)
  +        if ctx.p == 1:
  +            ctx.noise.fill_(0)
  +        else:
  +            ctx.noise.bernoulli_(1 - ctx.p).div_(1 - ctx.p)
  +        ctx.noise = ctx.noise.expand_as(input)
  +        output.mul_(ctx.noise)
   
           return output
  
verification_info: {"golden_diff": "diff --git a/torch/nn/_functions/dropout.py b/torch/nn/_functions/dropout.py\n--- a/torch/nn/_functions/dropout.py\n+++ b/torch/nn/_functions/dropout.py\n@@ -25,20 +25,22 @@\n         ctx.train = train\n         ctx.inplace = inplace\n \n+        if ctx.p == 0 or not ctx.train:\n+            return input\n+\n         if ctx.inplace:\n             ctx.mark_dirty(input)\n             output = input\n         else:\n             output = input.clone()\n \n-        if ctx.p > 0 and ctx.train:\n-            ctx.noise = cls._make_noise(input)\n-            if ctx.p == 1:\n-                ctx.noise.fill_(0)\n-            else:\n-                ctx.noise.bernoulli_(1 - ctx.p).div_(1 - ctx.p)\n-            ctx.noise = ctx.noise.expand_as(input)\n-            output.mul_(ctx.noise)\n+        ctx.noise = cls._make_noise(input)\n+        if ctx.p == 1:\n+            ctx.noise.fill_(0)\n+        else:\n+            ctx.noise.bernoulli_(1 - ctx.p).div_(1 - ctx.p)\n+        ctx.noise = ctx.noise.expand_as(input)\n+        output.mul_(ctx.noise)\n \n         return output\n", "issue": "JIT leaks memory in nets with dropout layer \nif I add a dropout layer in FC layers, a JIT compiled forward pass of a net starts to seriously leak memory in the backward pass. I made a simple test case below. When it runs, you will see a rapidly increasing GPU/CPU memory usage in nvidia-smi/top and it will crash with an out of memory exception soon. Without the dropout layer or the backward pass, it works fine. \r\n\r\n```python\r\nimport torch\r\nfrom torch import jit\r\nimport torch.nn as nn\r\nfrom torch.autograd import Variable\r\n\r\nclass TestNet(nn.Module):\r\n    def __init__(self):\r\n        super(TestNet, self).__init__()\r\n        self.net1 = nn.Linear(100, 200)\r\n        self.net2 = nn.Linear(200, 1)\r\n        self.sigmoid = nn.Sigmoid()\r\n        self.ReLU = nn.ReLU(inplace=False)\r\n        self.drop = nn.Dropout(0.5)\r\n               \r\n    def forward(self, V):\r\n        return self.sigmoid(self.net2(self.drop(self.ReLU(self.net1(V))))).squeeze() \r\n\r\n\r\nuse_cuda = True\r\nnet = TestNet()\r\ncriterion = nn.BCELoss()\r\nif use_cuda:\r\n    net.cuda()\r\n    criterion.cuda()\r\n    V = Variable(torch.randn(100, 100)).cuda()\r\n    label = Variable(torch.randn(100)).cuda()\r\nelse:\r\n    V = Variable(torch.randn(100, 100))\r\n    label = Variable(torch.randn(100))\r\n\r\nnet.train()\r\nfwd = jit.compile(net.forward)\r\nfor i in range(0,1000000):\r\n    r = fwd(V)\r\n    err = criterion(r, label)\r\n    err.backward()   \r\n\r\n```\r\n\r\n\n", "before_files": [{"content": "import torch\nfrom torch.autograd.function import InplaceFunction\nfrom torch.autograd import Variable\nfrom itertools import repeat\n\n\nclass Dropout(InplaceFunction):\n\n    @staticmethod\n    def _make_noise(input):\n        return input.new().resize_as_(input)\n\n    @staticmethod\n    def symbolic(g, input, p=0.5, train=False, inplace=False):\n        # See Note [Export inplace]\n        r, _ = g.op(\"Dropout\", input, ratio_f=p, is_test_i=not train, outputs=2)\n        return r\n\n    @classmethod\n    def forward(cls, ctx, input, p=0.5, train=False, inplace=False):\n        if p < 0 or p > 1:\n            raise ValueError(\"dropout probability has to be between 0 and 1, \"\n                             \"but got {}\".format(p))\n        ctx.p = p\n        ctx.train = train\n        ctx.inplace = inplace\n\n        if ctx.inplace:\n            ctx.mark_dirty(input)\n            output = input\n        else:\n            output = input.clone()\n\n        if ctx.p > 0 and ctx.train:\n            ctx.noise = cls._make_noise(input)\n            if ctx.p == 1:\n                ctx.noise.fill_(0)\n            else:\n                ctx.noise.bernoulli_(1 - ctx.p).div_(1 - ctx.p)\n            ctx.noise = ctx.noise.expand_as(input)\n            output.mul_(ctx.noise)\n\n        return output\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        if ctx.p > 0 and ctx.train:\n            return grad_output.mul(Variable(ctx.noise)), None, None, None\n        else:\n            return grad_output, None, None, None\n\n\nclass FeatureDropout(Dropout):\n\n    @staticmethod\n    def symbolic(g, input, p=0.5, train=False, inplace=False):\n        # See Note [Export inplace]\n        # NB: In inference mode, FeatureDropout is exported as an identity op.\n        from torch.onnx.symbolic import _unimplemented\n        if train:\n            return _unimplemented(\"FeatureDropout\", \"training mode\")\n        return input\n\n    @staticmethod\n    def _make_noise(input):\n        return input.new().resize_(input.size(0), input.size(1),\n                                   *repeat(1, input.dim() - 2))\n", "path": "torch/nn/_functions/dropout.py"}], "after_files": [{"content": "import torch\nfrom torch.autograd.function import InplaceFunction\nfrom torch.autograd import Variable\nfrom itertools import repeat\n\n\nclass Dropout(InplaceFunction):\n\n    @staticmethod\n    def _make_noise(input):\n        return input.new().resize_as_(input)\n\n    @staticmethod\n    def symbolic(g, input, p=0.5, train=False, inplace=False):\n        # See Note [Export inplace]\n        r, _ = g.op(\"Dropout\", input, ratio_f=p, is_test_i=not train, outputs=2)\n        return r\n\n    @classmethod\n    def forward(cls, ctx, input, p=0.5, train=False, inplace=False):\n        if p < 0 or p > 1:\n            raise ValueError(\"dropout probability has to be between 0 and 1, \"\n                             \"but got {}\".format(p))\n        ctx.p = p\n        ctx.train = train\n        ctx.inplace = inplace\n\n        if ctx.p == 0 or not ctx.train:\n            return input\n\n        if ctx.inplace:\n            ctx.mark_dirty(input)\n            output = input\n        else:\n            output = input.clone()\n\n        ctx.noise = cls._make_noise(input)\n        if ctx.p == 1:\n            ctx.noise.fill_(0)\n        else:\n            ctx.noise.bernoulli_(1 - ctx.p).div_(1 - ctx.p)\n        ctx.noise = ctx.noise.expand_as(input)\n        output.mul_(ctx.noise)\n\n        return output\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        if ctx.p > 0 and ctx.train:\n            return grad_output.mul(Variable(ctx.noise)), None, None, None\n        else:\n            return grad_output, None, None, None\n\n\nclass FeatureDropout(Dropout):\n\n    @staticmethod\n    def symbolic(g, input, p=0.5, train=False, inplace=False):\n        # See Note [Export inplace]\n        # NB: In inference mode, FeatureDropout is exported as an identity op.\n        from torch.onnx.symbolic import _unimplemented\n        if train:\n            return _unimplemented(\"FeatureDropout\", \"training mode\")\n        return input\n\n    @staticmethod\n    def _make_noise(input):\n        return input.new().resize_(input.size(0), input.size(1),\n                                   *repeat(1, input.dim() - 2))\n", "path": "torch/nn/_functions/dropout.py"}]}