Update app.py

app.py

@@ -21,7 +21,6 @@ data_normalized = scaler.fit_transform(data.reshape(-1, 1))
 
 seq_length = 4
 
-
 class LSTM(nn.Module):
     def __init__(self, input_size=INPUT_SIZE, hidden_layer_size=HIDDEN_LAYER_SIZE, output_size=OUTPUT_SIZE):
         super(LSTM, self).__init__()
@@ -36,14 +35,17 @@ class LSTM(nn.Module):
         predictions = self.linear(lstm_out.view(len(input_seq), -1))
         return predictions[-1]
 
-with open('LSTM_MODEL.pkl', 'rb') as f:
-    model = pickle.load(f)
-
 def prepare_custom_input(last_values, seq_length, scaler):
     last_values_normalized = scaler.transform(np.array(last_values).reshape(-1, 1))
     input_seq = torch.from_numpy(last_values_normalized).float()
     return input_seq.view(-1)
 
+model_path = 'lstm_model.pth'
+
+model = LSTMModel(INPUT_SIZE, HIDDEN_LAYER_SIZE, OUTPUT_SIZE)
+
+model.load_state_dict(torch.load(model_path))
+
 model.eval()
 
 def predict_and_plot(week1, week2, week3, week4):