From 7fba32127b8c8f25c2e12b03d1abda0866b3ebfb Mon Sep 17 00:00:00 2001 From: jjanezhang Date: Mon, 13 May 2024 14:22:41 -0700 Subject: [PATCH] removing rich install --- scripts/train/train.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/scripts/train/train.py b/scripts/train/train.py index 880d4f2350..4ba901597b 100644 --- a/scripts/train/train.py +++ b/scripts/train/train.py @@ -21,20 +21,16 @@ from composer.utils import dist, get_device, reproducibility from omegaconf import DictConfig from omegaconf import OmegaConf as om -from rich.traceback import install +from llmfoundry.callbacks import AsyncEval +from llmfoundry.data.dataloader import build_dataloader from llmfoundry.eval.metrics.nlp import InContextLearningMetric +from llmfoundry.layers_registry import ffns_with_megablocks from llmfoundry.utils import ( find_mosaicml_logger, log_train_analytics, maybe_create_mosaicml_logger, ) - -install() - -from llmfoundry.callbacks import AsyncEval -from llmfoundry.data.dataloader import build_dataloader -from llmfoundry.layers_registry import ffns_with_megablocks from llmfoundry.utils.builders import ( add_metrics_to_eval_loaders, build_algorithm,