-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtrain.py
30 lines (21 loc) · 848 Bytes
/
train.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import dotenv
import hydra
from omegaconf import DictConfig
# load environment variables from `.env` file if it exists
# recursively searches for `.env` in all folders starting from work dir
dotenv.load_dotenv(override=True)
# import torch; torch.set_anomaly_enabled(True)
# import torch; torch.use_deterministic_algorithms(True)
# import torch; torch.backends.cudnn.benchmark = False
@hydra.main(config_path="configs/", config_name="train.yaml", version_base="1.1")
def main(config: DictConfig):
# Imports can be nested inside @hydra.main to optimize tab completion
# https://github.com/facebookresearch/hydra/issues/934
from src import utils
from src.training_pipeline import train
# Applies optional utilities
utils.extras(config)
# Train model
return train(config)
if __name__ == "__main__":
main()