text: from typing import List
  
  try:
      import torch
      import torch.nn as nn
      from torch.utils.data import DataLoader
      from torch.optim.optimizer import Optimizer
  except ModuleNotFoundError:
      raise ModuleNotFoundError("Pytorch not found. Install with pip install torch")
  
  from timeatlas.abstract import AbstractBaseModel
  
  
  class LSTMPrediction(nn.Module, AbstractBaseModel):
  
      def __init__(self, n_features, n_hidden, seq_len, optimizer: Optimizer, loss_function, n_layers=1,
              bias=True, batch_first=False, dropout=0, bidirectional=False, horizon: int = 1, **kwargs):
          """
  
          Args:
              model: object instance of a LSTM from timeatlas (e.g. model=Prediction(<parameters>)
              epochs: number of epochs on the training set
              optimizer: class of the optimizer -> see Pytroch -> torch.optim
              loss_function: a function that calculates the loss -> see PyTorch -> torch.nn
              n_features: Number of features as input
              n_hidden: number of cells in each hidden layer
              seq_len: length of each featurs
              n_layers: number of hidden layer
              bias: If False, then the layer does not use bias weights b_ih and b_hh.
              batch_first: If True, then the input and output tensors are provided as (batch, seq, feature).
              dropout: If non-zero, introduces a Dropout layer on the outputs of each LSTM layer except the last layer, with dropout probability equal to dropout.
              bidirectional: if True, becomes a bidirectional LSTM.
              horizon: the number of returned values from the Linear output layer
              **kwargs: kwargs for the optimizer
          """
  
          super(LSTMPrediction, self).__init__()
  
          self.n_hidden = n_hidden
          self.seq_len = seq_len
          self.n_layers = n_layers
          self.bias = bias
          self.batch_first = batch_first
          self.dropout = dropout
          self.bidirectional = bidirectional
          self.out_features = horizon
  
          self.lstm = nn.LSTM(input_size=n_features,
                              hidden_size=n_hidden,
                              num_layers=n_layers,
                              bias=bias,
                              batch_first=batch_first,
                              dropout=dropout,
                              bidirectional=bidirectional
                              )
  
          self.linear = nn.Linear(in_features=n_hidden, out_features=self.out_features)
          self.reset_hidden_state()
  
          # optimizer and loss function for the fit
          self.optimizer = optimizer(self.parameters(), **kwargs)
          self.loss_function = loss_function
  
          # saving the history
          self.loss = []
  
      def reset_hidden_state(self):
          self.hidden = (
              torch.zeros(self.n_layers, self.seq_len, self.n_hidden).double(),
              torch.zeros(self.n_layers, self.seq_len, self.n_hidden).double()
          )
  
      def forward(self, sequences):
          lstm_out, self.hidden = self.lstm(sequences.view(len(sequences), self.seq_len, -1),
                                            self.hidden)
  
          last_time_step = lstm_out.view(self.seq_len, len(sequences), self.n_hidden)[-1]
          y_pred = self.linear(last_time_step)
          return y_pred
  
      def fit(self, series: DataLoader, epochs: int = 10, verbose: bool = True):
          super().fit(series)
          for i in range(epochs):
              for seq, labels in self.X_train:
  
                  self.reset_hidden_state()
  
                  self.optimizer.zero_grad()
  
                  y_pred = self(seq)
  
                  single_loss = self.loss_function(y_pred.squeeze(), labels)
                  single_loss.backward()
                  self.optimizer.step()
  
                  self.loss.append(single_loss.item())
  
          if verbose:
              print(f'epoch: {i:3} loss: {single_loss.item():10.8f}')
  
      def predict(self, horizon: int = 1, testset: DataLoader = None) -> List:
          super().predict(horizon)
  
          if testset is None:
              raise Exception(f'the testset has to be of type {type(DataLoader)}, got {type(testset)}')
  
          self.eval()
          predictions = []
          for seq, l in testset:
              y_pred = self(seq)
              predictions.append(y_pred)
  
          return predictions
  