# Buggy code
def objective(trial: optuna.Trial) -> float:
    # Filenames for each trial must be made unique in order to access each checkpoint.
    checkpoint_callback = pl.callbacks.ModelCheckpoint(
        os.path.join(model_path, "trial_{}".format(trial.number), "{epoch}"),
        monitor="val_loss",
    )

    # The default logger in PyTorch Lightning writes to event files to be consumed by
    # TensorBoard. We don't use any logger here as it requires us to implement several abstract
    # methods. Instead we setup a simple callback, that saves metrics from each validation step.
    metrics_callback = MetricsCallback()
    learning_rate_callback = LearningRateMonitor()
    logger = TensorBoardLogger(log_dir, name="optuna", version=trial.number)
    gradient_clip_val = trial.suggest_loguniform(
        "gradient_clip_val", *gradient_clip_val_range
    )
    trainer_kwargs.setdefault("gpus", [0] if torch.cuda.is_available() else None)
    trainer = pl.Trainer(
        checkpoint_callback=checkpoint_callback,
        max_epochs=max_epochs,
        gradient_clip_val=gradient_clip_val,
        callbacks=[
            metrics_callback,
            learning_rate_callback,
            PyTorchLightningPruningCallback(trial, monitor="val_loss"),
        ],
        logger=logger,
        progress_bar_refresh_rate=[0, 1][optuna_verbose < optuna.logging.INFO],
        weights_summary=[None, "top"][optuna_verbose < optuna.logging.INFO],
        **trainer_kwargs,
    )

    # create model
    hidden_size = trial.suggest_int("hidden_size", *hidden_size_range, log=True)
    kwargs["loss"] = copy.deepcopy(loss)
    model = TemporalFusionTransformer.from_dataset(
        train_dataloader.dataset,
        dropout=trial.suggest_uniform("dropout", *dropout_range),
        hidden_size=hidden_size,
        hidden_continuous_size=trial.suggest_int(
            "hidden_continuous_size",
            hidden_continuous_size_range[0],
            min(hidden_continuous_size_range[1], hidden_size),
            log=True,
        ),
        attention_head_size=trial.suggest_int(
            "attention_head_size", *attention_head_size_range
        ),
        log_interval=-1,
        **kwargs,
    )
    # find good learning rate
    if use_learning_rate_finder:
        lr_trainer = pl.Trainer(
            gradient_clip_val=gradient_clip_val,
            gpus=[0] if torch.cuda.is_available() else None,
            logger=False,
            progress_bar_refresh_rate=0,
            weights_summary=None,
        )
        res = lr_trainer.tuner.lr_find(
            model,
            train_dataloader=train_dataloader,
            val_dataloaders=val_dataloader,
            early_stop_threshold=10000,
            min_lr=learning_rate_range[0],
            num_training=100,
            max_lr=learning_rate_range[1],
        )

        loss_finite = np.isfinite(res.results["loss"])
        lr_smoothed, loss_smoothed = sm.nonparametric.lowess(
            np.asarray(res.results["loss"])[loss_finite],
            np.asarray(res.results["lr"])[loss_finite],
            frac=1.0 / 10.0,
        )[10:-1].T
        optimal_idx = np.gradient(loss_smoothed).argmin()
        optimal_lr = lr_smoothed[optimal_idx]
        optuna_logger.info(f"Using learning rate of {optimal_lr:.3g}")
        # add learning rate artificially
        model.hparams.learning_rate = trial.suggest_uniform(
            "learning_rate", optimal_lr, optimal_lr
        )
    else:
        model.hparams.learning_rate = trial.suggest_loguniform(
            "learning_rate", *learning_rate_range
        )

    # fit
    trainer.fit(
        model, train_dataloader=train_dataloader, val_dataloaders=val_dataloader
    )

    # report result
    return metrics_callback.metrics[-1]["val_loss"].item()


# 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 objective(trial: optuna.Trial) -> float:
    # Filenames for each trial must be made unique in order to access each checkpoint.
    checkpoint_callback = pl.callbacks.ModelCheckpoint(
        os.path.join(model_path, "trial_{}".format(trial.number), "{epoch}"),
        monitor="val_loss",
    )

    # The default logger in PyTorch Lightning writes to event files to be consumed by
    # TensorBoard. We don't use any logger here as it requires us to implement several abstract
    # methods. Instead we setup a simple callback, that saves metrics from each validation step.
    metrics_callback = MetricsCallback()
    learning_rate_callback = LearningRateMonitor()
    logger = TensorBoardLogger(log_dir, name="optuna", version=trial.number)
    gradient_clip_val = trial.suggest_loguniform(
        "gradient_clip_val", *gradient_clip_val_range
    )
    trainer_kwargs.setdefault("gpus", [0] if torch.cuda.is_available() else None)
    trainer = pl.Trainer(
        checkpoint_callback=checkpoint_callback,
        max_epochs=max_epochs,
        gradient_clip_val=gradient_clip_val,
        callbacks=[
            metrics_callback,
            learning_rate_callback,
            PyTorchLightningPruningCallback(trial, monitor="val_loss"),
        ],
        logger=logger,
        progress_bar_refresh_rate=[0, 1][optuna_verbose < optuna.logging.INFO],
        weights_summary=[None, "top"][optuna_verbose < optuna.logging.INFO],
        **trainer_kwargs,
    )

    # create model
    hidden_size = trial.suggest_int("hidden_size", *hidden_size_range, log=True)
    kwargs["loss"] = copy.deepcopy(loss)
    model = TemporalFusionTransformer.from_dataset(
        train_dataloader.dataset,
        dropout=trial.suggest_uniform("dropout", *dropout_range),
        hidden_size=hidden_size,
        hidden_continuous_size=trial.suggest_int(
            "hidden_continuous_size",
            hidden_continuous_size_range[0],
            min(hidden_continuous_size_range[1], hidden_size),
            log=True,
        ),
        attention_head_size=trial.suggest_int(
            "attention_head_size", *attention_head_size_range
        ),
        log_interval=-1,
        **kwargs,
    )
    # find good learning rate
    if use_learning_rate_finder:
        lr_trainer = pl.Trainer(
            gradient_clip_val=gradient_clip_val,
            gpus=[0] if torch.cuda.is_available() else None,
            logger=False,
            progress_bar_refresh_rate=0,
            weights_summary=None,
        )
        res = lr_trainer.tuner.lr_find(
            model,
            train_dataloader=train_dataloader,
            val_dataloaders=val_dataloader,
            early_stop_threshold=10000,
            min_lr=learning_rate_range[0],
            num_training=100,
            max_lr=learning_rate_range[1],
        )

        loss_finite = np.isfinite(res.results["loss"])
        if (
            loss_finite.sum() > 3
        ):  # at least 3 valid values required for learning rate finder
            lr_smoothed, loss_smoothed = sm.nonparametric.lowess(
                np.asarray(res.results["loss"])[loss_finite],
                np.asarray(res.results["lr"])[loss_finite],
                frac=1.0 / 10.0,
            )[min(loss_finite.sum() - 3, 10) : -1].T
            optimal_idx = np.gradient(loss_smoothed).argmin()
            optimal_lr = lr_smoothed[optimal_idx]
        else:
            optimal_idx = np.asarray(res.results["loss"]).argmin()
            optimal_lr = res.results["lr"][optimal_idx]
        optuna_logger.info(f"Using learning rate of {optimal_lr:.3g}")
        # add learning rate artificially
        model.hparams.learning_rate = trial.suggest_uniform(
            "learning_rate", optimal_lr, optimal_lr
        )
    else:
        model.hparams.learning_rate = trial.suggest_loguniform(
            "learning_rate", *learning_rate_range
        )

    # fit
    trainer.fit(
        model, train_dataloader=train_dataloader, val_dataloaders=val_dataloader
    )

    # report result
    return metrics_callback.metrics[-1]["val_loss"].item()