NASA RUL project - get lstm results

get_lstm_results.py

from tqdm import tqdm

def get_predictions_and_labels_lstm(dataset):
  predictions = []
  labels = []
  for item in tqdm(test_dataset):
    X, y = item
    _, output = trained_model(X.unsqueeze(dim=0))
    predictions.append(output.item())
    labels.append(y.item())
  return predictions, labels

print(f'RMSE on train set: {mean_squared_error(*get_predictions_and_labels_lstm(data_module.train_dataset), squared=False)}')
print(f'RMSE on test set: {mean_squared_error(*get_predictions_and_labels_lstm(data_module.test_dataset), squared=False)}')