chore: linting

Signed-off-by: Brandon Groth <[email protected]>

Author: BrandonGroth

tests/models/conftest.py

@@ -22,6 +22,7 @@
 import os
 
 # Third Party
+from torch.utils.data import DataLoader, TensorDataset
 from torchvision.io import read_image
 from torchvision.models import ResNet50_Weights, ViT_B_16_Weights, resnet50, vit_b_16
 from transformers import (
@@ -1346,7 +1347,7 @@ def input_tiny() -> DataLoader:
         dataset,
         batch_size=batch_size,
         shuffle=False,
-        collate_fn=lambda batch: tuple(torch.stack(samples) for samples in zip(*batch))
+        collate_fn=lambda batch: tuple(torch.stack(samples) for samples in zip(*batch)),
     )
 
 

tests/models/test_model_utils.py

@@ -238,4 +238,3 @@ def check_linear_dtypes(state_dict: dict, linear_names: list):
         for k, v in state_dict.items()
         if all(n not in k for n in linear_names) or not k.endswith(".weight")
     )
-

tests/models/test_save_aiu.py

@@ -3,11 +3,7 @@
 import pytest
 
 # Local
-from .test_model_utils import (
-    check_linear_dtypes,
-    delete_config,
-    load_state_dict,
-)
+from .test_model_utils import check_linear_dtypes, delete_config, load_state_dict
 from fms_mo import qmodel_prep
 from fms_mo.utils.aiu_utils import save_for_aiu
 
@@ -61,32 +57,33 @@ def test_large_outlier_bert(
         qcfg_bert (dict): Fake tiny input
         bert_linear_names (list): Quantized config for Bert
     """
+    # Third Party
     import torch
 
     # Break every tensor channel with a large magnitude outlier
-    for k,v in model_tiny_bert.state_dict().items():
+    for k, v in model_tiny_bert.state_dict().items():
         if k.endswith(".weight") and any(n in k for n in bert_linear_names):
-            v[:,0] = 1.21
+            v[:, 0] = 1.21
 
     # Set recomputation for narrow weights and prep
     qcfg_bert["recompute_narrow_weights"] = True
     qmodel_prep(model_tiny_bert, input_tiny, qcfg_bert, use_dynamo=True)
 
     # Qmax should break the quantization with an outlier to have skinny distribution
     layer2stdev: dict[str, torch.Tensor] = {}
-    for k,v in model_tiny_bert.state_dict().items():
+    for k, v in model_tiny_bert.state_dict().items():
         if k.endswith(".weight") and any(n in k for n in bert_linear_names):
             layer2stdev[k] = v.to(torch.float32).std(dim=-1)
 
     save_for_aiu(model_tiny_bert, qcfg=qcfg_bert, verbose=True)
     state_dict = load_state_dict()
 
     # Loaded model w/ recomputed SAWB should have widened channel quantization stdev
-    for k,v in state_dict.items():
+    for k, v in state_dict.items():
         if k.endswith(".weight") and any(n in k for n in bert_linear_names):
             perCh_stdev_model = layer2stdev.get(k)
             perCh_stdev_loaded = v.to(torch.float32).std(dim=-1)
-           
+
             assert torch.all(perCh_stdev_loaded >= perCh_stdev_model)
 
 
@@ -136,4 +133,3 @@ def test_save_model_granite(
     # Fetch saved state dict
     state_dict = load_state_dict()
     check_linear_dtypes(state_dict, granite_linear_names)
-