models/base.py

import pytorch_lightning as pl
import torch


class BaseLightningModule(pl.LightningModule):
    """
    Minimal LightningModule base for the demo.
    """

    def configure_optimizers(self):

        lr = getattr(self.hparams, "lr", 1e-3)
        return torch.optim.Adam(self.parameters(), lr=lr)