From d549499d76a7e4fe44bf2898654c28d4f51f483e Mon Sep 17 00:00:00 2001 From: v-chen_data Date: Thu, 6 Jun 2024 12:51:02 -0700 Subject: [PATCH] move distributed --- tests/models/test_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/test_model.py b/tests/models/test_model.py index a62a7dd114..78f7ee0d5c 100644 --- a/tests/models/test_model.py +++ b/tests/models/test_model.py @@ -15,7 +15,7 @@ from composer.core.precision import Precision, get_precision_context from composer.models.huggingface import maybe_get_underlying_model from composer.optim import DecoupledAdamW -from composer.trainer.dist_strategy import prepare_fsdp_module +from composer.distributed.dist_strategy import prepare_fsdp_module from composer.utils import dist, get_device, reproducibility from omegaconf import DictConfig, ListConfig from omegaconf import OmegaConf as om