# Buggy code
def distribution(
    self,
    distr_args,
    loc: Optional[Tensor] = None,
    scale: Optional[Tensor] = None,
) -> NegativeBinomial:
    assert loc is None
    mu, alpha = distr_args
    if scale is None:
        return NegativeBinomial(mu, alpha)
    else:
        F = getF(mu)
        mu = F.broadcast_mul(mu, scale)
        alpha = F.broadcast_mul(alpha, F.sqrt(scale + 1.0))
        return NegativeBinomial(mu, alpha, F)


# Traceback (AssertionError)
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-8-920902ee83c0> in <module>
22 )
23
---> 24 predictor = estimator.train(train_dataset)

/opt/amazon/lib/python3.6/site-packages/gluonts/model/estimator.py in train(self, training_data, validation_data)
221         self, training_data: Dataset, validation_data: Optional[Dataset] = None
222     ) -> Predictor:
--> 223         return self.train_model(training_data, validation_data).predictor

/opt/amazon/lib/python3.6/site-packages/gluonts/model/estimator.py in train_model(self, training_data, validation_data)
206             input_names=get_hybrid_forward_input_names(trained_net),
207             train_iter=training_data_loader,
--> 208             validation_iter=validation_data_loader,
209         )
210

/opt/amazon/lib/python3.6/site-packages/gluonts/trainer/_base.py in __call__(self, net, input_names, train_iter, validation_iter)
295                     )
296
--> 297                     epoch_loss = loop(epoch_no, train_iter)
298                     if is_validation_available:
299                         epoch_loss = loop(

/opt/amazon/lib/python3.6/site-packages/gluonts/trainer/_base.py in loop(epoch_no, batch_iter, is_training)
235
236                             with mx.autograd.record():
--> 237                                 output = net(*inputs)
238
239                                 # network can returns several outputs, the first being always the loss

/opt/amazon/lib/python3.6/site-packages/mxnet/gluon/block.py in __call__(self, *args)
546             hook(self, args)
547
--> 548         out = self.forward(*args)
549
550         for hook in self._forward_hooks.values():

/opt/amazon/lib/python3.6/site-packages/mxnet/gluon/block.py in forward(self, x, *args)
923                     params = {i: j.data(ctx) for i, j in self._reg_params.items()}
924
--> 925                 return self.hybrid_forward(ndarray, x, *args, **params)
926
927         assert isinstance(x, Symbol), \

/opt/amazon/lib/python3.6/site-packages/gluonts/model/deepar/_network.py in hybrid_forward(self, F, feat_static_cat, feat_static_real, past_time_feat, past_target, past_observed_values, future_time_feat, future_target, future_observed_values)
380             future_time_feat=future_time_feat,
381             future_target=future_target,
--> 382             future_observed_values=future_observed_values,
383         )
384

/opt/amazon/lib/python3.6/site-packages/gluonts/model/deepar/_network.py in distribution(self, feat_static_cat, feat_static_real, past_time_feat, past_target, past_observed_values, future_time_feat, future_target, future_observed_values)
336         distr_args = self.proj_distr_args(rnn_outputs)
337
--> 338         return self.distr_output.distribution(distr_args, scale=scale)
339
340     # noinspection PyMethodOverriding,PyPep8Naming

/opt/amazon/lib/python3.6/site-packages/gluonts/distribution/dirichlet_multinomial.py in distribution(self, distr_args, loc, scale)
182
183     def distribution(self, distr_args, loc=None, scale=None) -> Distribution:
--> 184         assert loc is None and scale is None
185         distr = DirichletMultinomial(self.dim, self.n_trials, distr_args)
186         return distr

AssertionError:

# Fixed code
def distribution(
    self,
    distr_args,
    loc: Optional[Tensor] = None,
    scale: Optional[Tensor] = None,
) -> NegativeBinomial:
    mu, alpha = distr_args
    if scale is None:
        return NegativeBinomial(mu, alpha)
    else:
        F = getF(mu)
        mu = F.broadcast_mul(mu, scale)
        alpha = F.broadcast_mul(alpha, F.sqrt(scale + 1.0))
        return NegativeBinomial(mu, alpha, F)