We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Support "device" keyword argument
Adds the optional "device" keyword argument to Tensor and Storage
constructors and .new methods.


--- 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/distributions/categorical.py`
Content:
```
1 import torch
2 from torch.autograd import Variable
3 from torch.distributions import constraints
4 from torch.distributions.distribution import Distribution
5 from torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property
6 
7 
8 class Categorical(Distribution):
9     r"""
10     Creates a categorical distribution parameterized by `probs`.
11 
12     .. note::
13         It is equivalent to the distribution that ``multinomial()`` samples from.
14 
15     Samples are integers from `0 ... K-1` where `K` is probs.size(-1).
16 
17     If `probs` is 1D with length-`K`, each element is the relative probability
18     of sampling the class at that index.
19 
20     If `probs` is 2D, it is treated as a batch of probability vectors.
21 
22     See also: :func:`torch.multinomial`
23 
24     Example::
25 
26         >>> m = Categorical(torch.Tensor([ 0.25, 0.25, 0.25, 0.25 ]))
27         >>> m.sample()  # equal probability of 0, 1, 2, 3
28          3
29         [torch.LongTensor of size 1]
30 
31     Args:
32         probs (Tensor or Variable): event probabilities
33     """
34     params = {'probs': constraints.simplex}
35     has_enumerate_support = True
36 
37     def __init__(self, probs=None, logits=None):
38         if (probs is None) == (logits is None):
39             raise ValueError("Either `probs` or `logits` must be specified, but not both.")
40         if probs is not None:
41             self.probs = probs / probs.sum(-1, keepdim=True)
42         else:
43             self.logits = logits - log_sum_exp(logits)
44         batch_shape = self.probs.size()[:-1] if probs is not None else self.logits.size()[:-1]
45         super(Categorical, self).__init__(batch_shape)
46 
47     @constraints.dependent_property
48     def support(self):
49         return constraints.integer_interval(0, self.probs.size()[-1] - 1)
50 
51     @lazy_property
52     def logits(self):
53         return probs_to_logits(self.probs)
54 
55     @lazy_property
56     def probs(self):
57         return logits_to_probs(self.logits)
58 
59     def sample(self, sample_shape=torch.Size()):
60         num_events = self.probs.size()[-1]
61         sample_shape = self._extended_shape(sample_shape)
62         param_shape = sample_shape + self.probs.size()[-1:]
63         probs = self.probs.expand(param_shape)
64         probs_2d = probs.contiguous().view(-1, num_events)
65         sample_2d = torch.multinomial(probs_2d, 1, True)
66         return sample_2d.contiguous().view(sample_shape)
67 
68     def log_prob(self, value):
69         self._validate_log_prob_arg(value)
70         param_shape = value.size() + self.probs.size()[-1:]
71         log_pmf = self.logits.expand(param_shape)
72         return log_pmf.gather(-1, value.unsqueeze(-1).long()).squeeze(-1)
73 
74     def entropy(self):
75         p_log_p = self.logits * self.probs
76         return -p_log_p.sum(-1)
77 
78     def enumerate_support(self):
79         num_events = self.probs.size()[-1]
80         values = torch.arange(num_events).long()
81         values = values.view((-1,) + (1,) * len(self._batch_shape))
82         values = values.expand((-1,) + self._batch_shape)
83         if self.probs.is_cuda:
84             values = values.cuda(self.probs.get_device())
85         if isinstance(self.probs, Variable):
86             values = Variable(values)
87         return values
88
```
--- 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/torch/distributions/categorical.py b/torch/distributions/categorical.py
--- a/torch/distributions/categorical.py
+++ b/torch/distributions/categorical.py
@@ -2,7 +2,7 @@
 from torch.autograd import Variable
 from torch.distributions import constraints
 from torch.distributions.distribution import Distribution
-from torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property
+from torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property, broadcast_all
 
 
 class Categorical(Distribution):
@@ -41,6 +41,7 @@
             self.probs = probs / probs.sum(-1, keepdim=True)
         else:
             self.logits = logits - log_sum_exp(logits)
+        self._num_events = self.probs.size()[-1] if self.probs is not None else self.logits.size()[-1]
         batch_shape = self.probs.size()[:-1] if probs is not None else self.logits.size()[:-1]
         super(Categorical, self).__init__(batch_shape)
 
@@ -67,7 +68,9 @@
 
     def log_prob(self, value):
         self._validate_log_prob_arg(value)
-        param_shape = value.size() + self.probs.size()[-1:]
+        value_shape = torch._C._infer_size(value.size(), self.batch_shape) if self.batch_shape else value.size()
+        param_shape = value_shape + (self._num_events,)
+        value = value.expand(value_shape)
         log_pmf = self.logits.expand(param_shape)
         return log_pmf.gather(-1, value.unsqueeze(-1).long()).squeeze(-1)

{"golden_diff": "diff --git a/torch/distributions/categorical.py b/torch/distributions/categorical.py\n--- a/torch/distributions/categorical.py\n+++ b/torch/distributions/categorical.py\n@@ -2,7 +2,7 @@\n from torch.autograd import Variable\n from torch.distributions import constraints\n from torch.distributions.distribution import Distribution\n-from torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property\n+from torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property, broadcast_all\n \n \n class Categorical(Distribution):\n@@ -41,6 +41,7 @@\n             self.probs = probs / probs.sum(-1, keepdim=True)\n         else:\n             self.logits = logits - log_sum_exp(logits)\n+        self._num_events = self.probs.size()[-1] if self.probs is not None else self.logits.size()[-1]\n         batch_shape = self.probs.size()[:-1] if probs is not None else self.logits.size()[:-1]\n         super(Categorical, self).__init__(batch_shape)\n \n@@ -67,7 +68,9 @@\n \n     def log_prob(self, value):\n         self._validate_log_prob_arg(value)\n-        param_shape = value.size() + self.probs.size()[-1:]\n+        value_shape = torch._C._infer_size(value.size(), self.batch_shape) if self.batch_shape else value.size()\n+        param_shape = value_shape + (self._num_events,)\n+        value = value.expand(value_shape)\n         log_pmf = self.logits.expand(param_shape)\n         return log_pmf.gather(-1, value.unsqueeze(-1).long()).squeeze(-1)\n", "issue": "Support \"device\" keyword argument\nAdds the optional \"device\" keyword argument to Tensor and Storage\nconstructors and .new methods.\n\n", "before_files": [{"content": "import torch\nfrom torch.autograd import Variable\nfrom torch.distributions import constraints\nfrom torch.distributions.distribution import Distribution\nfrom torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property\n\n\nclass Categorical(Distribution):\n    r\"\"\"\n    Creates a categorical distribution parameterized by `probs`.\n\n    .. note::\n        It is equivalent to the distribution that ``multinomial()`` samples from.\n\n    Samples are integers from `0 ... K-1` where `K` is probs.size(-1).\n\n    If `probs` is 1D with length-`K`, each element is the relative probability\n    of sampling the class at that index.\n\n    If `probs` is 2D, it is treated as a batch of probability vectors.\n\n    See also: :func:`torch.multinomial`\n\n    Example::\n\n        >>> m = Categorical(torch.Tensor([ 0.25, 0.25, 0.25, 0.25 ]))\n        >>> m.sample()  # equal probability of 0, 1, 2, 3\n         3\n        [torch.LongTensor of size 1]\n\n    Args:\n        probs (Tensor or Variable): event probabilities\n    \"\"\"\n    params = {'probs': constraints.simplex}\n    has_enumerate_support = True\n\n    def __init__(self, probs=None, logits=None):\n        if (probs is None) == (logits is None):\n            raise ValueError(\"Either `probs` or `logits` must be specified, but not both.\")\n        if probs is not None:\n            self.probs = probs / probs.sum(-1, keepdim=True)\n        else:\n            self.logits = logits - log_sum_exp(logits)\n        batch_shape = self.probs.size()[:-1] if probs is not None else self.logits.size()[:-1]\n        super(Categorical, self).__init__(batch_shape)\n\n    @constraints.dependent_property\n    def support(self):\n        return constraints.integer_interval(0, self.probs.size()[-1] - 1)\n\n    @lazy_property\n    def logits(self):\n        return probs_to_logits(self.probs)\n\n    @lazy_property\n    def probs(self):\n        return logits_to_probs(self.logits)\n\n    def sample(self, sample_shape=torch.Size()):\n        num_events = self.probs.size()[-1]\n        sample_shape = self._extended_shape(sample_shape)\n        param_shape = sample_shape + self.probs.size()[-1:]\n        probs = self.probs.expand(param_shape)\n        probs_2d = probs.contiguous().view(-1, num_events)\n        sample_2d = torch.multinomial(probs_2d, 1, True)\n        return sample_2d.contiguous().view(sample_shape)\n\n    def log_prob(self, value):\n        self._validate_log_prob_arg(value)\n        param_shape = value.size() + self.probs.size()[-1:]\n        log_pmf = self.logits.expand(param_shape)\n        return log_pmf.gather(-1, value.unsqueeze(-1).long()).squeeze(-1)\n\n    def entropy(self):\n        p_log_p = self.logits * self.probs\n        return -p_log_p.sum(-1)\n\n    def enumerate_support(self):\n        num_events = self.probs.size()[-1]\n        values = torch.arange(num_events).long()\n        values = values.view((-1,) + (1,) * len(self._batch_shape))\n        values = values.expand((-1,) + self._batch_shape)\n        if self.probs.is_cuda:\n            values = values.cuda(self.probs.get_device())\n        if isinstance(self.probs, Variable):\n            values = Variable(values)\n        return values\n", "path": "torch/distributions/categorical.py"}], "after_files": [{"content": "import torch\nfrom torch.autograd import Variable\nfrom torch.distributions import constraints\nfrom torch.distributions.distribution import Distribution\nfrom torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property, broadcast_all\n\n\nclass Categorical(Distribution):\n    r\"\"\"\n    Creates a categorical distribution parameterized by `probs`.\n\n    .. note::\n        It is equivalent to the distribution that ``multinomial()`` samples from.\n\n    Samples are integers from `0 ... K-1` where `K` is probs.size(-1).\n\n    If `probs` is 1D with length-`K`, each element is the relative probability\n    of sampling the class at that index.\n\n    If `probs` is 2D, it is treated as a batch of probability vectors.\n\n    See also: :func:`torch.multinomial`\n\n    Example::\n\n        >>> m = Categorical(torch.Tensor([ 0.25, 0.25, 0.25, 0.25 ]))\n        >>> m.sample()  # equal probability of 0, 1, 2, 3\n         3\n        [torch.LongTensor of size 1]\n\n    Args:\n        probs (Tensor or Variable): event probabilities\n    \"\"\"\n    params = {'probs': constraints.simplex}\n    has_enumerate_support = True\n\n    def __init__(self, probs=None, logits=None):\n        if (probs is None) == (logits is None):\n            raise ValueError(\"Either `probs` or `logits` must be specified, but not both.\")\n        if probs is not None:\n            self.probs = probs / probs.sum(-1, keepdim=True)\n        else:\n            self.logits = logits - log_sum_exp(logits)\n        self._num_events = self.probs.size()[-1] if self.probs is not None else self.logits.size()[-1]\n        batch_shape = self.probs.size()[:-1] if probs is not None else self.logits.size()[:-1]\n        super(Categorical, self).__init__(batch_shape)\n\n    @constraints.dependent_property\n    def support(self):\n        return constraints.integer_interval(0, self.probs.size()[-1] - 1)\n\n    @lazy_property\n    def logits(self):\n        return probs_to_logits(self.probs)\n\n    @lazy_property\n    def probs(self):\n        return logits_to_probs(self.logits)\n\n    def sample(self, sample_shape=torch.Size()):\n        num_events = self.probs.size()[-1]\n        sample_shape = self._extended_shape(sample_shape)\n        param_shape = sample_shape + self.probs.size()[-1:]\n        probs = self.probs.expand(param_shape)\n        probs_2d = probs.contiguous().view(-1, num_events)\n        sample_2d = torch.multinomial(probs_2d, 1, True)\n        return sample_2d.contiguous().view(sample_shape)\n\n    def log_prob(self, value):\n        self._validate_log_prob_arg(value)\n        value_shape = torch._C._infer_size(value.size(), self.batch_shape) if self.batch_shape else value.size()\n        param_shape = value_shape + (self._num_events,)\n        value = value.expand(value_shape)\n        log_pmf = self.logits.expand(param_shape)\n        return log_pmf.gather(-1, value.unsqueeze(-1).long()).squeeze(-1)\n\n    def entropy(self):\n        p_log_p = self.logits * self.probs\n        return -p_log_p.sum(-1)\n\n    def enumerate_support(self):\n        num_events = self.probs.size()[-1]\n        values = torch.arange(num_events).long()\n        values = values.view((-1,) + (1,) * len(self._batch_shape))\n        values = values.expand((-1,) + self._batch_shape)\n        if self.probs.is_cuda:\n            values = values.cuda(self.probs.get_device())\n        if isinstance(self.probs, Variable):\n            values = Variable(values)\n        return values\n", "path": "torch/distributions/categorical.py"}]}