input: 
  import numpy as np
  
  class LinearRegressionModel:
      def __init__(self, input_dim: int, output_dim: int):
          """ Initialises the weights and bias of the linear regression model
          Arguments:
              input_dim {int} -- Number of input features
              output_dim {int} -- Number of outputs
          """
          self.weights = np.zeros((input_dim, output_dim))
          self.bias = np.zeros((1, output_dim))
  
      def forward(self, x: np.ndarray) -> np.ndarray:
          """ Calculates the forward pass of the linear regression model
          Arguments:
              x {numpy array} -- Input data of size (N, input_dim)
          Returns:
              out {numpy array} -- Output of the linear regression model of size (N, output_dim)
          """
          out = np.dot(x, self.weights) + self.bias
          return out
  
  # Example usage:
  # >>> model = LinearRegressionModel(input_dim=3, output_dim=1)
  # >>> model.forward(np.array([[1, 2, 3]]))
  # array([[0.]])
  
output:     
      def backward(self, x: np.ndarray, grad_output: np.ndarray) -> np.ndarray:
          """ Calculates the backward pass of the linear regression model
          Arguments:
              x {numpy array} -- Input data of size (N, input_dim)
              grad_output {numpy array} -- Gradient of the loss with respect to the output of size (N, output_dim)
          Returns:
              grad_input {numpy array} -- Gradient of the loss with respect to the input of size (N, input_dim)
          """
          grad_input = np.dot(grad_output, self.weights.T)
          self.weights -= np.dot(x.T, grad_output)
          self.bias -= np.sum(grad_output, axis=0, keepdims=True)
          return grad_input
  
      def train(self, x_train: np.ndarray, y_train: np.ndarray, learning_rate: float, epochs: int):
          """ Trains the linear regression model
          Arguments:
              x_train {numpy array} -- Training input data of size (N, input_dim)
              y_train {numpy array} -- Training output data of size (N, output_dim)
              learning_rate {float} -- Learning rate for gradient descent
              epochs {int} -- Number of epochs to train
          """
          for epoch in range(epochs):
              y_pred = self.forward(x_train)
              loss = np.mean((y_pred - y_train) ** 2)
              grad_output = 2 * (y_pred - y_train) / x_train.shape[0]
              self.backward(x_train, grad_output)
              print(f"Epoch {epoch + 1}/{epochs}, Loss: {loss}")
  
  # Example usage:
  # >>> model = LinearRegressionModel(input_dim=3, output_dim=1)
  # >>> x_train = np.array([[1, 2, 3], [4, 5, 6]])
  # >>> y_train = np.array([[1], [2]])
  # >>> model.train(x_train, y_train, learning_rate=0.01, epochs=100)
  