diff --git a/requirements.txt b/requirements.txt index 5fef6d4..13ff2e0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,4 @@ torch==1.11.0 numpy~=1.22.3 matplotlib==3.5.2 sacred==0.8.2 +pymongo==4.1.1 \ No newline at end of file diff --git a/train_model.py b/train_model.py index fe2a5f3..c82006d 100644 --- a/train_model.py +++ b/train_model.py @@ -3,7 +3,7 @@ import argparse import numpy as np import pandas as pd import torch -from sacred.observers import FileStorageObserver +from sacred.observers import FileStorageObserver, MongoObserver from torch import nn from torch.utils.data import DataLoader, Dataset from sacred import Experiment @@ -49,6 +49,8 @@ def main(batch_size, epochs, _run): def setup_experiment(): ex = Experiment('Predict power output for a given time') ex.observers.append(FileStorageObserver('sacred_runs')) + # ex.observers.append(MongoObserver(url='mongodb://mongo_user:mongo_password_IUM_2022@localhost:27017', + # db_name='sacred')) return ex