We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Getting NaN loss after some time with DeepAR and NegativeBinomial
## Description
I have a dataset with 129k time series of demand for products within some branches of a retailer. I'm having trouble with NaN loss ever since I started using GluonTS. At first, I discovered some NaN and Infinity values within my FEAT_DYNAMIC_REAL. But even after solving those issues, I'm still getting NaN loss after some time. With the new release (which included some improvements on the NegativeBinomial, I'm getting a much lower loss value and it's taking longer for the problem to occur, but it still occurs.

I decided to put a conditional breakpoint in the log_prob() to stop whenever a NaN value is generated and used hybridize=False. It's super strange, but the error didn't happen yet (It's far past the point that it usually happens). I'll wait some more time, but it might be the case that it only happens with hybridize=True? If that's the case, it'd be very odd and problematic, since it runs much slower with hybridize=False.

## To Reproduce

Unfortunately, I can't provide the dataset to reproduce the problem. But here are some information:
- Number of Time Series: 128.976
- TARGET: Count values, with a lot of zeros
- FEAT_STATIC_CAT: I'm using categorical features that identify the product (its id, its category, subcategory, and so on) and the branch (its id, its city, state, and so on).
- FEAT_DYNAMIC_REAL: I'm using holidays in one-hot encoding, a binary identifying if there's a special price and the relative discount (calculated using the discount and the price).
- batch_size: 32
- num_batches_per_epoch: 10000
- learning_rate: 1e-4
- clip_gradient: 1.0
- patience: 10
- minimum_learning_rate: 5e-5
- weight_decay: 1e-8
- num_layers: 2
- num_cells: 100
- embedding_dimension: [2, 7, 20, 50, 10, 30, 60, 80, 200]
- context_length: 90
- scaling: True
- freq: 1D

```python
put code here
```

## Error message or code output
Here is one example where it happened in the first epoch. But sometimes it happened later.

```
Traceback (most recent call last):
  File "/opt/anaconda3/envs/gazin/lib/python3.6/site-packages/luigi/worker.py", line 199, in run
    new_deps = self._run_get_new_deps()
  File "/opt/anaconda3/envs/gazin/lib/python3.6/site-packages/luigi/worker.py", line 141, in _run_get_new_deps
    task_gen = self.task.run()
  File "/root/work/log_inteligente_ai/gazin/task/model/base.py", line 164, in run
    dict(fit=self.fit, cross_validation=self.cross_validate)[self.training_type]()
  File "/root/work/log_inteligente_ai/gazin/task/model/gluonts.py", line 425, in fit
    predictor = self.train_model()
  File "/root/work/log_inteligente_ai/gazin/task/model/gluonts.py", line 416, in train_model
    train_iter=training_data_loader,
  File "/opt/anaconda3/envs/gazin/lib/python3.6/site-packages/gluonts/trainer/_base.py", line 328, in __call__
    "Got NaN in first epoch. Try reducing initial learning rate."
gluonts.core.exception.GluonTSUserError: Got NaN in first epoch. Try reducing initial learning rate.
```

## Environment
- Operating system: Linux
- Python version: 3.6
- GluonTS version: 0.5.0



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/gluonts/mx/distribution/neg_binomial.py`
Content:
```
1 # Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License").
4 # You may not use this file except in compliance with the License.
5 # A copy of the License is located at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # or in the "license" file accompanying this file. This file is distributed
10 # on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
11 # express or implied. See the License for the specific language governing
12 # permissions and limitations under the License.
13 
14 from typing import Dict, List, Optional, Tuple
15 
16 import numpy as np
17 
18 from gluonts.core.component import validated
19 from gluonts.mx import Tensor
20 
21 from .deterministic import DeterministicOutput
22 from .distribution import Distribution, _sample_multiple, getF, softplus
23 from .distribution_output import DistributionOutput
24 from .mixture import MixtureDistributionOutput
25 
26 
27 class NegativeBinomial(Distribution):
28     r"""
29     Negative binomial distribution, i.e. the distribution of the number of
30     successes in a sequence of independent Bernoulli trials.
31 
32     Parameters
33     ----------
34     mu
35         Tensor containing the means, of shape `(*batch_shape, *event_shape)`.
36     alpha
37         Tensor of the shape parameters, of shape `(*batch_shape, *event_shape)`.
38     F
39     """
40 
41     is_reparameterizable = False
42 
43     @validated()
44     def __init__(self, mu: Tensor, alpha: Tensor) -> None:
45         self.mu = mu
46         self.alpha = alpha
47 
48     @property
49     def F(self):
50         return getF(self.mu)
51 
52     @property
53     def batch_shape(self) -> Tuple:
54         return self.mu.shape
55 
56     @property
57     def event_shape(self) -> Tuple:
58         return ()
59 
60     @property
61     def event_dim(self) -> int:
62         return 0
63 
64     def log_prob(self, x: Tensor) -> Tensor:
65         alphaInv = 1.0 / self.alpha
66         alpha_times_mu = self.alpha * self.mu
67         F = self.F
68         ll = (
69             x * F.log(alpha_times_mu / (1.0 + alpha_times_mu))
70             - alphaInv * F.log1p(alpha_times_mu)
71             + F.gammaln(x + alphaInv)
72             - F.gammaln(x + 1.0)
73             - F.gammaln(alphaInv)
74         )
75         return ll
76 
77     @property
78     def mean(self) -> Tensor:
79         return self.mu
80 
81     @property
82     def stddev(self) -> Tensor:
83         return self.F.sqrt(self.mu * (1.0 + self.mu * self.alpha))
84 
85     def sample(
86         self, num_samples: Optional[int] = None, dtype=np.float32
87     ) -> Tensor:
88         def s(mu: Tensor, alpha: Tensor) -> Tensor:
89             F = self.F
90             r = 1.0 / alpha
91             theta = alpha * mu
92             return F.random.poisson(lam=F.random.gamma(r, theta), dtype=dtype)
93 
94         return _sample_multiple(
95             s, mu=self.mu, alpha=self.alpha, num_samples=num_samples
96         )
97 
98     @property
99     def args(self) -> List:
100         return [self.mu, self.alpha]
101 
102 
103 class NegativeBinomialOutput(DistributionOutput):
104     args_dim: Dict[str, int] = {"mu": 1, "alpha": 1}
105     distr_cls: type = NegativeBinomial
106 
107     @classmethod
108     def domain_map(cls, F, mu, alpha):
109         epsilon = np.finfo(cls._dtype).eps  # machine epsilon
110 
111         mu = softplus(F, mu) + epsilon
112         alpha = softplus(F, alpha) + epsilon
113         return mu.squeeze(axis=-1), alpha.squeeze(axis=-1)
114 
115     # Overwrites the parent class method.
116     # We cannot scale using the affine transformation since negative binomial should return integers.
117     # Instead we scale the parameters.
118     def distribution(
119         self,
120         distr_args,
121         loc: Optional[Tensor] = None,
122         scale: Optional[Tensor] = None,
123     ) -> NegativeBinomial:
124         mu, alpha = distr_args
125         if scale is None:
126             return NegativeBinomial(mu, alpha)
127         else:
128             F = getF(mu)
129             mu = F.broadcast_mul(mu, scale)
130             return NegativeBinomial(mu, alpha, F)
131 
132     @property
133     def event_shape(self) -> Tuple:
134         return ()
135 
136 
137 def ZeroInflatedNegativeBinomialOutput() -> MixtureDistributionOutput:
138     return MixtureDistributionOutput(
139         distr_outputs=[NegativeBinomialOutput(), DeterministicOutput(0)]
140     )
141
```
--- 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/src/gluonts/mx/distribution/neg_binomial.py b/src/gluonts/mx/distribution/neg_binomial.py
--- a/src/gluonts/mx/distribution/neg_binomial.py
+++ b/src/gluonts/mx/distribution/neg_binomial.py
@@ -107,9 +107,8 @@
     @classmethod
     def domain_map(cls, F, mu, alpha):
         epsilon = np.finfo(cls._dtype).eps  # machine epsilon
-
-        mu = softplus(F, mu) + epsilon
-        alpha = softplus(F, alpha) + epsilon
+        mu = F.maximum(softplus(F, mu), epsilon)
+        alpha = F.maximum(softplus(F, alpha), epsilon)
         return mu.squeeze(axis=-1), alpha.squeeze(axis=-1)
 
     # Overwrites the parent class method.

{"golden_diff": "diff --git a/src/gluonts/mx/distribution/neg_binomial.py b/src/gluonts/mx/distribution/neg_binomial.py\n--- a/src/gluonts/mx/distribution/neg_binomial.py\n+++ b/src/gluonts/mx/distribution/neg_binomial.py\n@@ -107,9 +107,8 @@\n     @classmethod\n     def domain_map(cls, F, mu, alpha):\n         epsilon = np.finfo(cls._dtype).eps  # machine epsilon\n-\n-        mu = softplus(F, mu) + epsilon\n-        alpha = softplus(F, alpha) + epsilon\n+        mu = F.maximum(softplus(F, mu), epsilon)\n+        alpha = F.maximum(softplus(F, alpha), epsilon)\n         return mu.squeeze(axis=-1), alpha.squeeze(axis=-1)\n \n     # Overwrites the parent class method.\n", "issue": "Getting NaN loss after some time with DeepAR and NegativeBinomial\n## Description\r\nI have a dataset with 129k time series of demand for products within some branches of a retailer. I'm having trouble with NaN loss ever since I started using GluonTS. At first, I discovered some NaN and Infinity values within my FEAT_DYNAMIC_REAL. But even after solving those issues, I'm still getting NaN loss after some time. With the new release (which included some improvements on the NegativeBinomial, I'm getting a much lower loss value and it's taking longer for the problem to occur, but it still occurs.\r\n\r\nI decided to put a conditional breakpoint in the log_prob() to stop whenever a NaN value is generated and used hybridize=False. It's super strange, but the error didn't happen yet (It's far past the point that it usually happens). I'll wait some more time, but it might be the case that it only happens with hybridize=True? If that's the case, it'd be very odd and problematic, since it runs much slower with hybridize=False.\r\n\r\n## To Reproduce\r\n\r\nUnfortunately, I can't provide the dataset to reproduce the problem. But here are some information:\r\n- Number of Time Series: 128.976\r\n- TARGET: Count values, with a lot of zeros\r\n- FEAT_STATIC_CAT: I'm using categorical features that identify the product (its id, its category, subcategory, and so on) and the branch (its id, its city, state, and so on).\r\n- FEAT_DYNAMIC_REAL: I'm using holidays in one-hot encoding, a binary identifying if there's a special price and the relative discount (calculated using the discount and the price).\r\n- batch_size: 32\r\n- num_batches_per_epoch: 10000\r\n- learning_rate: 1e-4\r\n- clip_gradient: 1.0\r\n- patience: 10\r\n- minimum_learning_rate: 5e-5\r\n- weight_decay: 1e-8\r\n- num_layers: 2\r\n- num_cells: 100\r\n- embedding_dimension: [2, 7, 20, 50, 10, 30, 60, 80, 200]\r\n- context_length: 90\r\n- scaling: True\r\n- freq: 1D\r\n\r\n```python\r\nput code here\r\n```\r\n\r\n## Error message or code output\r\nHere is one example where it happened in the first epoch. But sometimes it happened later.\r\n\r\n```\r\nTraceback (most recent call last):\r\n  File \"/opt/anaconda3/envs/gazin/lib/python3.6/site-packages/luigi/worker.py\", line 199, in run\r\n    new_deps = self._run_get_new_deps()\r\n  File \"/opt/anaconda3/envs/gazin/lib/python3.6/site-packages/luigi/worker.py\", line 141, in _run_get_new_deps\r\n    task_gen = self.task.run()\r\n  File \"/root/work/log_inteligente_ai/gazin/task/model/base.py\", line 164, in run\r\n    dict(fit=self.fit, cross_validation=self.cross_validate)[self.training_type]()\r\n  File \"/root/work/log_inteligente_ai/gazin/task/model/gluonts.py\", line 425, in fit\r\n    predictor = self.train_model()\r\n  File \"/root/work/log_inteligente_ai/gazin/task/model/gluonts.py\", line 416, in train_model\r\n    train_iter=training_data_loader,\r\n  File \"/opt/anaconda3/envs/gazin/lib/python3.6/site-packages/gluonts/trainer/_base.py\", line 328, in __call__\r\n    \"Got NaN in first epoch. Try reducing initial learning rate.\"\r\ngluonts.core.exception.GluonTSUserError: Got NaN in first epoch. Try reducing initial learning rate.\r\n```\r\n\r\n## Environment\r\n- Operating system: Linux\r\n- Python version: 3.6\r\n- GluonTS version: 0.5.0\r\n\r\n\n", "before_files": [{"content": "# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\").\n# You may not use this file except in compliance with the License.\n# A copy of the License is located at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# or in the \"license\" file accompanying this file. This file is distributed\n# on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either\n# express or implied. See the License for the specific language governing\n# permissions and limitations under the License.\n\nfrom typing import Dict, List, Optional, Tuple\n\nimport numpy as np\n\nfrom gluonts.core.component import validated\nfrom gluonts.mx import Tensor\n\nfrom .deterministic import DeterministicOutput\nfrom .distribution import Distribution, _sample_multiple, getF, softplus\nfrom .distribution_output import DistributionOutput\nfrom .mixture import MixtureDistributionOutput\n\n\nclass NegativeBinomial(Distribution):\n    r\"\"\"\n    Negative binomial distribution, i.e. the distribution of the number of\n    successes in a sequence of independent Bernoulli trials.\n\n    Parameters\n    ----------\n    mu\n        Tensor containing the means, of shape `(*batch_shape, *event_shape)`.\n    alpha\n        Tensor of the shape parameters, of shape `(*batch_shape, *event_shape)`.\n    F\n    \"\"\"\n\n    is_reparameterizable = False\n\n    @validated()\n    def __init__(self, mu: Tensor, alpha: Tensor) -> None:\n        self.mu = mu\n        self.alpha = alpha\n\n    @property\n    def F(self):\n        return getF(self.mu)\n\n    @property\n    def batch_shape(self) -> Tuple:\n        return self.mu.shape\n\n    @property\n    def event_shape(self) -> Tuple:\n        return ()\n\n    @property\n    def event_dim(self) -> int:\n        return 0\n\n    def log_prob(self, x: Tensor) -> Tensor:\n        alphaInv = 1.0 / self.alpha\n        alpha_times_mu = self.alpha * self.mu\n        F = self.F\n        ll = (\n            x * F.log(alpha_times_mu / (1.0 + alpha_times_mu))\n            - alphaInv * F.log1p(alpha_times_mu)\n            + F.gammaln(x + alphaInv)\n            - F.gammaln(x + 1.0)\n            - F.gammaln(alphaInv)\n        )\n        return ll\n\n    @property\n    def mean(self) -> Tensor:\n        return self.mu\n\n    @property\n    def stddev(self) -> Tensor:\n        return self.F.sqrt(self.mu * (1.0 + self.mu * self.alpha))\n\n    def sample(\n        self, num_samples: Optional[int] = None, dtype=np.float32\n    ) -> Tensor:\n        def s(mu: Tensor, alpha: Tensor) -> Tensor:\n            F = self.F\n            r = 1.0 / alpha\n            theta = alpha * mu\n            return F.random.poisson(lam=F.random.gamma(r, theta), dtype=dtype)\n\n        return _sample_multiple(\n            s, mu=self.mu, alpha=self.alpha, num_samples=num_samples\n        )\n\n    @property\n    def args(self) -> List:\n        return [self.mu, self.alpha]\n\n\nclass NegativeBinomialOutput(DistributionOutput):\n    args_dim: Dict[str, int] = {\"mu\": 1, \"alpha\": 1}\n    distr_cls: type = NegativeBinomial\n\n    @classmethod\n    def domain_map(cls, F, mu, alpha):\n        epsilon = np.finfo(cls._dtype).eps  # machine epsilon\n\n        mu = softplus(F, mu) + epsilon\n        alpha = softplus(F, alpha) + epsilon\n        return mu.squeeze(axis=-1), alpha.squeeze(axis=-1)\n\n    # Overwrites the parent class method.\n    # We cannot scale using the affine transformation since negative binomial should return integers.\n    # Instead we scale the parameters.\n    def distribution(\n        self,\n        distr_args,\n        loc: Optional[Tensor] = None,\n        scale: Optional[Tensor] = None,\n    ) -> NegativeBinomial:\n        mu, alpha = distr_args\n        if scale is None:\n            return NegativeBinomial(mu, alpha)\n        else:\n            F = getF(mu)\n            mu = F.broadcast_mul(mu, scale)\n            return NegativeBinomial(mu, alpha, F)\n\n    @property\n    def event_shape(self) -> Tuple:\n        return ()\n\n\ndef ZeroInflatedNegativeBinomialOutput() -> MixtureDistributionOutput:\n    return MixtureDistributionOutput(\n        distr_outputs=[NegativeBinomialOutput(), DeterministicOutput(0)]\n    )\n", "path": "src/gluonts/mx/distribution/neg_binomial.py"}], "after_files": [{"content": "# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\").\n# You may not use this file except in compliance with the License.\n# A copy of the License is located at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# or in the \"license\" file accompanying this file. This file is distributed\n# on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either\n# express or implied. See the License for the specific language governing\n# permissions and limitations under the License.\n\nfrom typing import Dict, List, Optional, Tuple\n\nimport numpy as np\n\nfrom gluonts.core.component import validated\nfrom gluonts.mx import Tensor\n\nfrom .deterministic import DeterministicOutput\nfrom .distribution import Distribution, _sample_multiple, getF, softplus\nfrom .distribution_output import DistributionOutput\nfrom .mixture import MixtureDistributionOutput\n\n\nclass NegativeBinomial(Distribution):\n    r\"\"\"\n    Negative binomial distribution, i.e. the distribution of the number of\n    successes in a sequence of independent Bernoulli trials.\n\n    Parameters\n    ----------\n    mu\n        Tensor containing the means, of shape `(*batch_shape, *event_shape)`.\n    alpha\n        Tensor of the shape parameters, of shape `(*batch_shape, *event_shape)`.\n    F\n    \"\"\"\n\n    is_reparameterizable = False\n\n    @validated()\n    def __init__(self, mu: Tensor, alpha: Tensor) -> None:\n        self.mu = mu\n        self.alpha = alpha\n\n    @property\n    def F(self):\n        return getF(self.mu)\n\n    @property\n    def batch_shape(self) -> Tuple:\n        return self.mu.shape\n\n    @property\n    def event_shape(self) -> Tuple:\n        return ()\n\n    @property\n    def event_dim(self) -> int:\n        return 0\n\n    def log_prob(self, x: Tensor) -> Tensor:\n        alphaInv = 1.0 / self.alpha\n        alpha_times_mu = self.alpha * self.mu\n        F = self.F\n        ll = (\n            x * F.log(alpha_times_mu / (1.0 + alpha_times_mu))\n            - alphaInv * F.log1p(alpha_times_mu)\n            + F.gammaln(x + alphaInv)\n            - F.gammaln(x + 1.0)\n            - F.gammaln(alphaInv)\n        )\n        return ll\n\n    @property\n    def mean(self) -> Tensor:\n        return self.mu\n\n    @property\n    def stddev(self) -> Tensor:\n        return self.F.sqrt(self.mu * (1.0 + self.mu * self.alpha))\n\n    def sample(\n        self, num_samples: Optional[int] = None, dtype=np.float32\n    ) -> Tensor:\n        def s(mu: Tensor, alpha: Tensor) -> Tensor:\n            F = self.F\n            r = 1.0 / alpha\n            theta = alpha * mu\n            return F.random.poisson(lam=F.random.gamma(r, theta), dtype=dtype)\n\n        return _sample_multiple(\n            s, mu=self.mu, alpha=self.alpha, num_samples=num_samples\n        )\n\n    @property\n    def args(self) -> List:\n        return [self.mu, self.alpha]\n\n\nclass NegativeBinomialOutput(DistributionOutput):\n    args_dim: Dict[str, int] = {\"mu\": 1, \"alpha\": 1}\n    distr_cls: type = NegativeBinomial\n\n    @classmethod\n    def domain_map(cls, F, mu, alpha):\n        epsilon = np.finfo(cls._dtype).eps  # machine epsilon\n        mu = F.maximum(softplus(F, mu), epsilon)\n        alpha = F.maximum(softplus(F, alpha), epsilon)\n        return mu.squeeze(axis=-1), alpha.squeeze(axis=-1)\n\n    # Overwrites the parent class method.\n    # We cannot scale using the affine transformation since negative binomial should return integers.\n    # Instead we scale the parameters.\n    def distribution(\n        self,\n        distr_args,\n        loc: Optional[Tensor] = None,\n        scale: Optional[Tensor] = None,\n    ) -> NegativeBinomial:\n        mu, alpha = distr_args\n        if scale is None:\n            return NegativeBinomial(mu, alpha)\n        else:\n            F = getF(mu)\n            mu = F.broadcast_mul(mu, scale)\n            return NegativeBinomial(mu, alpha, F)\n\n    @property\n    def event_shape(self) -> Tuple:\n        return ()\n\n\ndef ZeroInflatedNegativeBinomialOutput() -> MixtureDistributionOutput:\n    return MixtureDistributionOutput(\n        distr_outputs=[NegativeBinomialOutput(), DeterministicOutput(0)]\n    )\n", "path": "src/gluonts/mx/distribution/neg_binomial.py"}]}