# Buggy code
def rescale_parameters(
    self, parameters: torch.Tensor, target_scale: torch.Tensor, encoder: BaseEstimator
) -> torch.Tensor:
    assert encoder.transformation in ["logit"], (
        "Beta distribution is only compatible with logit transformation"
    )
    assert encoder.center, "Beta distribution requires normalizer to center data"

    scaled_mean = encoder(
        dict(prediction=parameters[..., 0], target_scale=target_scale)
    )
    # need to first transform target scale standard deviation in logit space to real space
    # we assume a normal distribution in logit space (we used a logit transform and a standard scaler)
    # and know that the variance of the beta distribution is limited by `scaled_mean * (1 - scaled_mean)`
    mean_derivative = scaled_mean * (1 - scaled_mean)

    # we can approximate variance as
    # torch.pow(torch.tanh(target_scale[..., 1].unsqueeze(1) * torch.sqrt(mean_derivative)), 2) * mean_derivative
    # shape is (positive) parameter * mean_derivative / var
    shape_scaler = torch.pow(
        torch.tanh(target_scale[..., 1].unsqueeze(1) * torch.sqrt(mean_derivative)), 2
    )
    scaled_shape = F.softplus(parameters[..., 1]) / shape_scaler
    return torch.stack([scaled_mean, scaled_shape], dim=-1)


# Traceback (RuntimeError)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-12-6644968460fc> in <module>
28 )
29
---> 30 trainer.fit(
31     deepar,
32     train_dataloader=train_dataloader,

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
468         self.call_hook('on_fit_start')
469
--> 470         results = self.accelerator_backend.train()
471         self.accelerator_backend.teardown()
472

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/accelerators/cpu_accelerator.py in train(self)
60
61         # train or test
---> 62         results = self.train_or_test()
63         return results
64

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py in train_or_test(self)
67             results = self.trainer.run_test()
68         else:
---> 69             results = self.trainer.train()
70         return results
71

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py in train(self)
519                 with self.profiler.profile("run_training_epoch"):
520                     # run train epoch
--> 521                     self.train_loop.run_training_epoch()
522
523                 if self.max_steps and self.max_steps <= self.global_step:

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py in run_training_epoch(self)
558             # ------------------------------------
559             with self.trainer.profiler.profile("run_training_batch"):
--> 560                 batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
561
562             # when returning -1 from train_step, we end epoch early

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py in run_training_batch(self, batch, batch_idx, dataloader_idx)
716
717                         # optimizer step
--> 718                         self.optimizer_step(optimizer, opt_idx, batch_idx, train_step_and_backward_closure)
719
720                     else:

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py in optimizer_step(self, optimizer, opt_idx, batch_idx, train_step_and_backward_closure)
491
492         # model hook
--> 493         model_ref.optimizer_step(
494             self.trainer.current_epoch,
495             batch_idx,

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py in optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure, on_tpu, using_native_amp, using_lbfgs)
1255             # wraps into LightingOptimizer only for running step
1256             optimizer = LightningOptimizer.to_lightning_optimizer(optimizer, self.trainer)
-> 1257         optimizer.step(closure=optimizer_closure)
1258
1259     def optimizer_zero_grad(

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py in step(self, closure, make_optimizer_step, *args, **kwargs)
276
277         if make_optimizer_step:
--> 278             self.__optimizer_step(*args, closure=closure, profiler_name=profiler_name, **kwargs)
279         else:
280             # make sure to call optimizer_closure when accumulating

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py in __optimizer_step(self, closure, profiler_name, *args, **kwargs)
134         else:
135             with trainer.profiler.profile(profiler_name):
--> 136                 optimizer.step(closure=closure, *args, **kwargs)
137
138         accelerator_backend = trainer.accelerator_backend

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_forecasting/optim.py in step(self, closure)
129             closure: A closure that reevaluates the model and returns the loss.
130         """
--> 131         _ = closure()
132         loss = None
133         # note - below is commented out b/c I have other work that passes back

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py in train_step_and_backward_closure()
706
707                         def train_step_and_backward_closure():
--> 708                             result = self.training_step_and_backward(
709                                 split_batch,
710                                 batch_idx,

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py in training_step_and_backward(self, split_batch, batch_idx, opt_idx, optimizer, hiddens)
814                 # backward pass
815                 with self.trainer.profiler.profile("model_backward"):
--> 816                     self.backward(result, optimizer, opt_idx)
817
818                 # hook - call this hook only

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py in backward(self, result, optimizer, opt_idx, *args, **kwargs)
834             self.trainer.accelerator_backend.backward(result, optimizer, opt_idx, *args, **kwargs)
835         else:
--> 836             result.closure_loss = self.trainer.accelerator_backend.backward(
837                 result.closure_loss, optimizer, opt_idx, *args, **kwargs
838             )

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py in backward(self, closure_loss, optimizer, opt_idx, *args, **kwargs)
102             # do backward pass
103             model = self.trainer.get_model()
--> 104             model.backward(closure_loss, optimizer, opt_idx, *args, **kwargs)
105
106             # once backward has been applied, release graph

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py in backward(self, loss, optimizer, optimizer_idx, *args, **kwargs)
1149         """
1150         if self.trainer.train_loop.automatic_optimization or self._running_manual_backward:
-> 1151             loss.backward(*args, **kwargs)
1152
1153     def toggle_optimizer(self, optimizer: Optimizer, optimizer_idx: int):

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
219                 retain_graph=retain_graph,
220                 create_graph=create_graph)
--> 221         torch.autograd.backward(self, gradient, retain_graph, create_graph)
222
223     def register_hook(self, hook):

~/miniconda3/envs/smartsilex-eda/lib/python3.8/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
128         retain_graph = create_graph
129
--> 130     Variable._execution_engine.run_backward(
131         tensors, grad_tensors_, retain_graph, create_graph,
132         allow_unreachable=True)  # allow_unreachable flag

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

# Fixed code
def rescale_parameters(
    self, parameters: torch.Tensor, target_scale: torch.Tensor, encoder: BaseEstimator
) -> torch.Tensor:
    assert encoder.transformation in ["logit"], (
        "Beta distribution is only compatible with logit transformation"
    )
    assert encoder.center, "Beta distribution requires normalizer to center data"

    scaled_mean = encoder(
        dict(prediction=parameters[..., 0], target_scale=target_scale)
    )
    # need to first transform target scale standard deviation in logit space to real space
    # we assume a normal distribution in logit space (we used a logit transform and a standard scaler)
    # and know that the variance of the beta distribution is limited by `scaled_mean * (1 - scaled_mean)`
    scaled_mean = (
        scaled_mean * (1 - 2 * self.eps) + self.eps
    )  # ensure that mean is not exactly 0 or 1
    mean_derivative = scaled_mean * (1 - scaled_mean)

    # we can approximate variance as
    # torch.pow(torch.tanh(target_scale[..., 1].unsqueeze(1) * torch.sqrt(mean_derivative)), 2) * mean_derivative
    # shape is (positive) parameter * mean_derivative / var
    shape_scaler = (
        torch.pow(
            torch.tanh(target_scale[..., 1].unsqueeze(1) * torch.sqrt(mean_derivative)),
            2,
        )
        + self.eps
    )
    scaled_shape = F.softplus(parameters[..., 1]) / shape_scaler
    return torch.stack([scaled_mean, scaled_shape], dim=-1)