prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  SyncBatchNorm doesn't support 2 dimensions input?
  Hi,
  I'm facing the issue that the program crash when the input for SyncBatchNorm is two dimensions. Here's the code:
  ```python
  import torch
  import apex
  
  model = apex.parallel.SyncBatchNorm(4).cuda()
  data = torch.rand((8,4)).cuda()
  output = model(data)
  ```
  When running the code, error raised like this:
  ```
  Traceback (most recent call last):
    File "syncbn_test.by", line 7, in <module>
      output = model(data)
    File "/usr/local/lib/python3.5/dist-packages/torch/nn/modules/module.py", line 489, in __call__
      result = self.forward(*input, **kwargs)
    File "/usr/local/lib/python3.5/dist-packages/apex/parallel/optimized_sync_batchnorm.py", line 81, in forward
      return SyncBatchnormFunction.apply(input, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, self.channel_last)
    File "/usr/local/lib/python3.5/dist-packages/apex/parallel/optimized_sync_batchnorm_kernel.py", line 27, in forward
      mean, var_biased = syncbn.welford_mean_var(input)
  RuntimeError: Dimension out of range (expected to be in range of [-2, 1], but got 2) (maybe_wrap_dim at /pytorch/aten/src/ATen/core/WrapDimMinimal.h:18)
  ```
  And everthing runs ok when `data` a 4 dims tensor. 
  
  Here is my environment:
  ```
  Ubuntu 16.04
  Python 3.5.2
  Pytorch 1.0.1, installed with "pip install torch"
  apex is installed with command:
    pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .
  cuda 10.0
  nvidia driver 410.72
  ```
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `apex/parallel/optimized_sync_batchnorm.py`
  Content:
  ```
  1 import torch
  2 from torch.nn.modules.batchnorm import _BatchNorm
  3 from torch.nn import functional as F
  4 
  5 import syncbn
  6 from .optimized_sync_batchnorm_kernel import SyncBatchnormFunction
  7 
  8 
  9 class SyncBatchNorm(_BatchNorm):
  10     """
  11     synchronized batch normalization module extented from `torch.nn.BatchNormNd`
  12     with the added stats reduction across multiple processes.
  13     :class:`apex.parallel.SyncBatchNorm` is designed to work with
  14     `DistributedDataParallel`.
  15 
  16     When running in training mode, the layer reduces stats across all processes
  17     to increase the effective batchsize for normalization layer. This is useful
  18     in applications where batch size is small on a given process that would
  19     diminish converged accuracy of the model. The model uses collective
  20     communication package from `torch.distributed`.
  21 
  22     When running in evaluation mode, the layer falls back to
  23     `torch.nn.functional.batch_norm`
  24 
  25     Args:
  26         num_features: :math:`C` from an expected input of size
  27             :math:`(N, C, L)` or :math:`L` from input of size :math:`(N, L)`
  28         eps: a value added to the denominator for numerical stability.
  29             Default: 1e-5
  30         momentum: the value used for the running_mean and running_var
  31             computation. Can be set to ``None`` for cumulative moving average
  32             (i.e. simple average). Default: 0.1
  33         affine: a boolean value that when set to ``True``, this module has
  34             learnable affine parameters. Default: ``True``
  35         track_running_stats: a boolean value that when set to ``True``, this
  36             module tracks the running mean and variance, and when set to ``False``,
  37             this module does not track such statistics and always uses batch
  38             statistics in both training and eval modes. Default: ``True``
  39         process_group: pass in a process group within which the stats of the
  40             mini-batch is being synchronized. ``None`` for using default process
  41             group
  42         channel_last: a boolean value that when set to ``True``, this module
  43             take the last dimension of the input tensor to be the channel
  44             dimension. Default: False
  45 
  46     Examples::
  47         >>> # channel first tensor
  48         >>> sbn = apex.parallel.SyncBatchNorm(100).cuda()
  49         >>> inp = torch.randn(10, 100, 14, 14).cuda()
  50         >>> out = sbn(inp)
  51         >>> inp = torch.randn(3, 100, 20).cuda()
  52         >>> out = sbn(inp)
  53         >>> # channel last tensor
  54         >>> sbn = apex.parallel.SyncBatchNorm(100, channel_last=True).cuda()
  55         >>> inp = torch.randn(10, 14, 14, 100).cuda()
  56     """
  57 
  58     def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last=False, fuse_relu=False):
  59         super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)
  60         self.process_group = process_group
  61         self.channel_last = channel_last
  62         self.fuse_relu = fuse_relu
  63 
  64     def _specify_process_group(self, process_group):
  65         self.process_group = process_group
  66 
  67     def _specify_channel_last(self, channel_last):
  68         self.channel_last = channel_last
  69 
  70     def forward(self, input, z = None):
  71         # if input.dim() == 2, we switch to channel_last for efficient memory accessing
  72         channel_last = self.channel_last if input.dim() != 2 else True
  73 
  74         if not self.training and self.track_running_stats and not self.channel_last and not self.fuse_relu and z == None:
  75             # fall back to pytorch implementation for inference
  76             return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)
  77         else:
  78             exponential_average_factor = 0.0
  79             if self.training and self.track_running_stats:
  80                 self.num_batches_tracked += 1
  81                 if self.momentum is None:
  82                     exponential_average_factor = 1.0 / float(self.num_batches_tracked)
  83                 else:
  84                     exponential_average_factor = self.momentum
  85             return SyncBatchnormFunction.apply(input, z, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, self.channel_last, self.fuse_relu)
  86
  ```
  --- 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/apex/parallel/optimized_sync_batchnorm.py b/apex/parallel/optimized_sync_batchnorm.py
  --- a/apex/parallel/optimized_sync_batchnorm.py
  +++ b/apex/parallel/optimized_sync_batchnorm.py
  @@ -71,7 +71,7 @@
           # if input.dim() == 2, we switch to channel_last for efficient memory accessing
           channel_last = self.channel_last if input.dim() != 2 else True
   
  -        if not self.training and self.track_running_stats and not self.channel_last and not self.fuse_relu and z == None:
  +        if not self.training and self.track_running_stats and not channel_last and not self.fuse_relu and z == None:
               # fall back to pytorch implementation for inference
               return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)
           else:
  @@ -82,4 +82,4 @@
                       exponential_average_factor = 1.0 / float(self.num_batches_tracked)
                   else:
                       exponential_average_factor = self.momentum
  -            return SyncBatchnormFunction.apply(input, z, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, self.channel_last, self.fuse_relu)
  +            return SyncBatchnormFunction.apply(input, z, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, channel_last, self.fuse_relu)
  
verification_info: {"golden_diff": "diff --git a/apex/parallel/optimized_sync_batchnorm.py b/apex/parallel/optimized_sync_batchnorm.py\n--- a/apex/parallel/optimized_sync_batchnorm.py\n+++ b/apex/parallel/optimized_sync_batchnorm.py\n@@ -71,7 +71,7 @@\n         # if input.dim() == 2, we switch to channel_last for efficient memory accessing\n         channel_last = self.channel_last if input.dim() != 2 else True\n \n-        if not self.training and self.track_running_stats and not self.channel_last and not self.fuse_relu and z == None:\n+        if not self.training and self.track_running_stats and not channel_last and not self.fuse_relu and z == None:\n             # fall back to pytorch implementation for inference\n             return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)\n         else:\n@@ -82,4 +82,4 @@\n                     exponential_average_factor = 1.0 / float(self.num_batches_tracked)\n                 else:\n                     exponential_average_factor = self.momentum\n-            return SyncBatchnormFunction.apply(input, z, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, self.channel_last, self.fuse_relu)\n+            return SyncBatchnormFunction.apply(input, z, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, channel_last, self.fuse_relu)\n", "issue": "SyncBatchNorm doesn't support 2 dimensions input?\nHi,\r\nI'm facing the issue that the program crash when the input for SyncBatchNorm is two dimensions. Here's the code:\r\n```python\r\nimport torch\r\nimport apex\r\n\r\nmodel = apex.parallel.SyncBatchNorm(4).cuda()\r\ndata = torch.rand((8,4)).cuda()\r\noutput = model(data)\r\n```\r\nWhen running the code, error raised like this:\r\n```\r\nTraceback (most recent call last):\r\n  File \"syncbn_test.by\", line 7, in <module>\r\n    output = model(data)\r\n  File \"/usr/local/lib/python3.5/dist-packages/torch/nn/modules/module.py\", line 489, in __call__\r\n    result = self.forward(*input, **kwargs)\r\n  File \"/usr/local/lib/python3.5/dist-packages/apex/parallel/optimized_sync_batchnorm.py\", line 81, in forward\r\n    return SyncBatchnormFunction.apply(input, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, self.channel_last)\r\n  File \"/usr/local/lib/python3.5/dist-packages/apex/parallel/optimized_sync_batchnorm_kernel.py\", line 27, in forward\r\n    mean, var_biased = syncbn.welford_mean_var(input)\r\nRuntimeError: Dimension out of range (expected to be in range of [-2, 1], but got 2) (maybe_wrap_dim at /pytorch/aten/src/ATen/core/WrapDimMinimal.h:18)\r\n```\r\nAnd everthing runs ok when `data` a 4 dims tensor. \r\n\r\nHere is my environment:\r\n```\r\nUbuntu 16.04\r\nPython 3.5.2\r\nPytorch 1.0.1, installed with \"pip install torch\"\r\napex is installed with command:\r\n  pip install -v --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" .\r\ncuda 10.0\r\nnvidia driver 410.72\r\n```\n", "before_files": [{"content": "import torch\nfrom torch.nn.modules.batchnorm import _BatchNorm\nfrom torch.nn import functional as F\n\nimport syncbn\nfrom .optimized_sync_batchnorm_kernel import SyncBatchnormFunction\n\n\nclass SyncBatchNorm(_BatchNorm):\n    \"\"\"\n    synchronized batch normalization module extented from `torch.nn.BatchNormNd`\n    with the added stats reduction across multiple processes.\n    :class:`apex.parallel.SyncBatchNorm` is designed to work with\n    `DistributedDataParallel`.\n\n    When running in training mode, the layer reduces stats across all processes\n    to increase the effective batchsize for normalization layer. This is useful\n    in applications where batch size is small on a given process that would\n    diminish converged accuracy of the model. The model uses collective\n    communication package from `torch.distributed`.\n\n    When running in evaluation mode, the layer falls back to\n    `torch.nn.functional.batch_norm`\n\n    Args:\n        num_features: :math:`C` from an expected input of size\n            :math:`(N, C, L)` or :math:`L` from input of size :math:`(N, L)`\n        eps: a value added to the denominator for numerical stability.\n            Default: 1e-5\n        momentum: the value used for the running_mean and running_var\n            computation. Can be set to ``None`` for cumulative moving average\n            (i.e. simple average). Default: 0.1\n        affine: a boolean value that when set to ``True``, this module has\n            learnable affine parameters. Default: ``True``\n        track_running_stats: a boolean value that when set to ``True``, this\n            module tracks the running mean and variance, and when set to ``False``,\n            this module does not track such statistics and always uses batch\n            statistics in both training and eval modes. Default: ``True``\n        process_group: pass in a process group within which the stats of the\n            mini-batch is being synchronized. ``None`` for using default process\n            group\n        channel_last: a boolean value that when set to ``True``, this module\n            take the last dimension of the input tensor to be the channel\n            dimension. Default: False\n\n    Examples::\n        >>> # channel first tensor\n        >>> sbn = apex.parallel.SyncBatchNorm(100).cuda()\n        >>> inp = torch.randn(10, 100, 14, 14).cuda()\n        >>> out = sbn(inp)\n        >>> inp = torch.randn(3, 100, 20).cuda()\n        >>> out = sbn(inp)\n        >>> # channel last tensor\n        >>> sbn = apex.parallel.SyncBatchNorm(100, channel_last=True).cuda()\n        >>> inp = torch.randn(10, 14, 14, 100).cuda()\n    \"\"\"\n\n    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last=False, fuse_relu=False):\n        super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)\n        self.process_group = process_group\n        self.channel_last = channel_last\n        self.fuse_relu = fuse_relu\n\n    def _specify_process_group(self, process_group):\n        self.process_group = process_group\n\n    def _specify_channel_last(self, channel_last):\n        self.channel_last = channel_last\n\n    def forward(self, input, z = None):\n        # if input.dim() == 2, we switch to channel_last for efficient memory accessing\n        channel_last = self.channel_last if input.dim() != 2 else True\n\n        if not self.training and self.track_running_stats and not self.channel_last and not self.fuse_relu and z == None:\n            # fall back to pytorch implementation for inference\n            return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)\n        else:\n            exponential_average_factor = 0.0\n            if self.training and self.track_running_stats:\n                self.num_batches_tracked += 1\n                if self.momentum is None:\n                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)\n                else:\n                    exponential_average_factor = self.momentum\n            return SyncBatchnormFunction.apply(input, z, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, self.channel_last, self.fuse_relu)\n", "path": "apex/parallel/optimized_sync_batchnorm.py"}], "after_files": [{"content": "import torch\nfrom torch.nn.modules.batchnorm import _BatchNorm\nfrom torch.nn import functional as F\n\nimport syncbn\nfrom .optimized_sync_batchnorm_kernel import SyncBatchnormFunction\n\n\nclass SyncBatchNorm(_BatchNorm):\n    \"\"\"\n    synchronized batch normalization module extented from `torch.nn.BatchNormNd`\n    with the added stats reduction across multiple processes.\n    :class:`apex.parallel.SyncBatchNorm` is designed to work with\n    `DistributedDataParallel`.\n\n    When running in training mode, the layer reduces stats across all processes\n    to increase the effective batchsize for normalization layer. This is useful\n    in applications where batch size is small on a given process that would\n    diminish converged accuracy of the model. The model uses collective\n    communication package from `torch.distributed`.\n\n    When running in evaluation mode, the layer falls back to\n    `torch.nn.functional.batch_norm`\n\n    Args:\n        num_features: :math:`C` from an expected input of size\n            :math:`(N, C, L)` or :math:`L` from input of size :math:`(N, L)`\n        eps: a value added to the denominator for numerical stability.\n            Default: 1e-5\n        momentum: the value used for the running_mean and running_var\n            computation. Can be set to ``None`` for cumulative moving average\n            (i.e. simple average). Default: 0.1\n        affine: a boolean value that when set to ``True``, this module has\n            learnable affine parameters. Default: ``True``\n        track_running_stats: a boolean value that when set to ``True``, this\n            module tracks the running mean and variance, and when set to ``False``,\n            this module does not track such statistics and always uses batch\n            statistics in both training and eval modes. Default: ``True``\n        process_group: pass in a process group within which the stats of the\n            mini-batch is being synchronized. ``None`` for using default process\n            group\n        channel_last: a boolean value that when set to ``True``, this module\n            take the last dimension of the input tensor to be the channel\n            dimension. Default: False\n\n    Examples::\n        >>> # channel first tensor\n        >>> sbn = apex.parallel.SyncBatchNorm(100).cuda()\n        >>> inp = torch.randn(10, 100, 14, 14).cuda()\n        >>> out = sbn(inp)\n        >>> inp = torch.randn(3, 100, 20).cuda()\n        >>> out = sbn(inp)\n        >>> # channel last tensor\n        >>> sbn = apex.parallel.SyncBatchNorm(100, channel_last=True).cuda()\n        >>> inp = torch.randn(10, 14, 14, 100).cuda()\n    \"\"\"\n\n    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last=False, fuse_relu=False):\n        super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)\n        self.process_group = process_group\n        self.channel_last = channel_last\n        self.fuse_relu = fuse_relu\n\n    def _specify_process_group(self, process_group):\n        self.process_group = process_group\n\n    def _specify_channel_last(self, channel_last):\n        self.channel_last = channel_last\n\n    def forward(self, input, z = None):\n        # if input.dim() == 2, we switch to channel_last for efficient memory accessing\n        channel_last = self.channel_last if input.dim() != 2 else True\n\n        if not self.training and self.track_running_stats and not channel_last and not self.fuse_relu and z == None:\n            # fall back to pytorch implementation for inference\n            return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)\n        else:\n            exponential_average_factor = 0.0\n            if self.training and self.track_running_stats:\n                self.num_batches_tracked += 1\n                if self.momentum is None:\n                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)\n                else:\n                    exponential_average_factor = self.momentum\n            return SyncBatchnormFunction.apply(input, z, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, channel_last, self.fuse_relu)\n", "path": "apex/parallel/optimized_sync_batchnorm.py"}]}