From e85fc95b6d3e4c8607061f2a44c7b3a51ac4bbb5 Mon Sep 17 00:00:00 2001
From: Irene Dea <deaairene@gmail.com>
Date: Mon, 23 Oct 2023 18:59:32 -0700
Subject: [PATCH] Fix code quality

---
 tests/fixtures/autouse.py        | 2 ++
 tests/fixtures/data.py           | 7 +++++--
 tests/test_flash_triton_torch.py | 1 -
 tests/test_hf_config.py          | 1 -
 tests/test_hf_v_mpt.py           | 3 ---
 tests/test_init_fn.py            | 1 -
 tests/test_llama_patch.py        | 1 -
 tests/test_model.py              | 2 +-
 tests/test_onnx.py               | 1 -
 9 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/tests/fixtures/autouse.py b/tests/fixtures/autouse.py
index 2aca76b5b8..c51ccfacb0 100644
--- a/tests/fixtures/autouse.py
+++ b/tests/fixtures/autouse.py
@@ -27,10 +27,12 @@ def clear_cuda_cache(request: pytest.FixtureRequest):
         torch.cuda.empty_cache()
         gc.collect()  # Only gc on GPU tests as it 2x slows down CPU tests
 
+
 @pytest.fixture
 def random_seed() -> int:
     return 17
 
+
 @pytest.fixture(autouse=True)
 def seed_all(random_seed: int):
     """Sets the seed for reproducibility."""
diff --git a/tests/fixtures/data.py b/tests/fixtures/data.py
index 3e7fb64d20..393bd420bd 100644
--- a/tests/fixtures/data.py
+++ b/tests/fixtures/data.py
@@ -47,8 +47,11 @@ def tiny_ft_dataloader(tiny_ft_dataset_path: str,
         'timeout': 0
     })
 
-    return build_finetuning_dataloader(
+    dataloader = build_finetuning_dataloader(
         dataloader_cfg,
         mpt_tokenizer,
         device_batch_size,
-    )
+    ).dataloader
+
+    assert isinstance(dataloader, DataLoader)
+    return dataloader
diff --git a/tests/test_flash_triton_torch.py b/tests/test_flash_triton_torch.py
index 37a4738bb5..e6fe8eb438 100644
--- a/tests/test_flash_triton_torch.py
+++ b/tests/test_flash_triton_torch.py
@@ -3,7 +3,6 @@
 
 import pytest
 import torch
-from composer.utils import reproducibility
 from omegaconf import OmegaConf as om
 
 
diff --git a/tests/test_hf_config.py b/tests/test_hf_config.py
index 5141516666..b47f267c55 100644
--- a/tests/test_hf_config.py
+++ b/tests/test_hf_config.py
@@ -9,7 +9,6 @@
 
 import pytest
 import torch
-from composer.utils import reproducibility
 from omegaconf import DictConfig
 from omegaconf import OmegaConf as om
 from transformers import AutoModelForCausalLM
diff --git a/tests/test_hf_v_mpt.py b/tests/test_hf_v_mpt.py
index 094059b329..46172faf35 100644
--- a/tests/test_hf_v_mpt.py
+++ b/tests/test_hf_v_mpt.py
@@ -5,7 +5,6 @@
 
 import pytest
 import torch
-from composer.utils import reproducibility
 from omegaconf import OmegaConf as om
 
 from llmfoundry import COMPOSER_MODEL_REGISTRY
@@ -52,8 +51,6 @@ def test_compare_hf_v_mpt(attn_impl: str, dropout: float, alibi: bool,
     batch_size = 2  # set batch size
     device = 'cuda'  # set decive
 
-
-
     # get hf gpt2 cfg
     hf_cfg = om.create({
         'model': {
diff --git a/tests/test_init_fn.py b/tests/test_init_fn.py
index 11bce99e47..6be2c5ca42 100644
--- a/tests/test_init_fn.py
+++ b/tests/test_init_fn.py
@@ -8,7 +8,6 @@
 
 import pytest
 import torch
-from composer.utils import reproducibility
 from omegaconf import DictConfig, ListConfig
 from omegaconf import OmegaConf as om
 from torch import nn
diff --git a/tests/test_llama_patch.py b/tests/test_llama_patch.py
index dec7057b1d..9b4c353d8a 100644
--- a/tests/test_llama_patch.py
+++ b/tests/test_llama_patch.py
@@ -6,7 +6,6 @@
 import pytest
 import torch
 import transformers
-from composer.utils import reproducibility
 from transformers.models.llama.modeling_llama import LlamaAttention
 
 from llmfoundry.models.layers.llama_attention_monkeypatch import (
diff --git a/tests/test_model.py b/tests/test_model.py
index 779226d0b6..67166bef68 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -16,7 +16,7 @@
 from composer.core.precision import Precision, get_precision_context
 from composer.optim import DecoupledAdamW
 from composer.trainer.dist_strategy import prepare_fsdp_module
-from composer.utils import dist, get_device, reproducibility
+from composer.utils import dist, get_device
 from omegaconf import DictConfig, ListConfig
 from omegaconf import OmegaConf as om
 from transformers import (AutoModelForCausalLM, AutoTokenizer, PreTrainedModel,
diff --git a/tests/test_onnx.py b/tests/test_onnx.py
index 603cba4276..d0e01746eb 100644
--- a/tests/test_onnx.py
+++ b/tests/test_onnx.py
@@ -4,7 +4,6 @@
 import pathlib
 
 import torch
-from composer.utils import reproducibility
 from transformers import AutoModelForCausalLM
 
 from llmfoundry import MPTConfig, MPTForCausalLM