remove icecream

mosaicml · Sep 25, 2024 · 1f025d8 · 1f025d8
1 parent 3a6dec6
commit 1f025d8
Show file tree

Hide file tree

Showing 2 changed files with 0 additions and 9 deletions.
diff --git a/llmfoundry/command_utils/train.py b/llmfoundry/command_utils/train.py
@@ -20,7 +20,6 @@
     cyclic_schedule,
 )
 from composer.utils import dist, get_device, reproducibility, ParallelismConfig, TPConfig, FSDPConfig
-from icecream import install, ic
 from omegaconf import DictConfig
 from omegaconf import OmegaConf as om
 

diff --git a/tests/models/utils/test_tp_strategy.py b/tests/models/utils/test_tp_strategy.py
@@ -1,18 +1,10 @@
-import pytest
-
-from composer.trainer.trainer import Trainer
-from composer.utils import dist
 from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel, PrepareModuleInput
 from torch.distributed._tensor import Replicate, Shard
-from torch.utils.data import DataLoader
 
 from llmfoundry.models.mpt.modeling_mpt import ComposerMPTCausalLM
 from llmfoundry.utils.builders import build_tp_strategy
 
 
-from icecream import install
-install()
-
 def test_ffn_tp_strategy_layer_plan():
 
     # Actual layer plan