From be935ba6f4c98ed3ad8afb999a44cc4c2a6a310e Mon Sep 17 00:00:00 2001
From: Ajay Dyavathi <ajaydyavathi@gmail.com>
Date: Wed, 11 Oct 2023 01:26:53 -0400
Subject: [PATCH] Update train.py

yaml.load(fp) deprecated, added a new way of loading yaml file.
---
 train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/train.py b/train.py
index eeabc578..f29cc80a 100644
--- a/train.py
+++ b/train.py
@@ -1,11 +1,11 @@
 import os
-import yaml
 import time
 import shutil
 import torch
 import random
 import argparse
 import numpy as np
+from yaml import load, Loader
 
 from torch.utils import data
 from tqdm import tqdm
@@ -214,7 +214,7 @@ def train(cfg, writer, logger):
     args = parser.parse_args()
 
     with open(args.config) as fp:
-        cfg = yaml.load(fp)
+        cfg = load(fp, Loader=Loader)
 
     run_id = random.randint(1, 100000)
     logdir = os.path.join("runs", os.path.basename(args.config)[:-4], str(run_id))