From 150016484ac4be3b34e873f6bc1e49ce2638c58f Mon Sep 17 00:00:00 2001
From: root <abhay.gupta@databricks.com>
Date: Mon, 29 Jul 2024 23:20:47 +0000
Subject: [PATCH] changes for pre-commit

---
 tests/models/layers/test_ffn.py | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/tests/models/layers/test_ffn.py b/tests/models/layers/test_ffn.py
index d6098bc80c..bb78763f58 100644
--- a/tests/models/layers/test_ffn.py
+++ b/tests/models/layers/test_ffn.py
@@ -3,11 +3,11 @@
 
 import pytest
 import torch
-import torch.nn as nn
 import torch.distributed as dist
+import torch.nn as nn
 
-from llmfoundry.models.layers.layer_builders import build_ffn
 from llmfoundry.models.layers.ffn import quickgelu_activation
+from llmfoundry.models.layers.layer_builders import build_ffn
 
 
 @pytest.mark.gpu
@@ -15,28 +15,29 @@ def test_quickgelu_activation():
     d_model = 32
     expansion_ratio = 1
     no_bias = True
-    ffn_config={
+    ffn_config = {
         'ffn_act_fn': {
             'name': 'quick_gelu',
         },
         'ffn_type': 'mptmlp',
     }
     rank: int = dist.get_rank()
-    device: torch.device = torch.device(f'cuda:{rank}')
+    device_str = f'cuda:{rank}'
+    device: torch.device = torch.device(device_str)
 
     ffn1 = build_ffn(
         name=ffn_config['ffn_type'],
         d_model=d_model,
         expansion_ratio=expansion_ratio,
-        device=device,
+        device=device_str,
         bias=not no_bias,
         ffn_kwargs=ffn_config,
     )
     assert (
         ffn1.act == quickgelu_activation
-    ), f"Expected quick_gelu activation function, got {ffn1.act}"
+    ), f'Expected quick_gelu activation function, got {ffn1.act}'
 
-    ffn_config={
+    ffn_config = {
         'ffn_act_fn': {
             'name': 'gelu',
         },
@@ -46,7 +47,7 @@ def test_quickgelu_activation():
         name=ffn_config['ffn_type'],
         d_model=d_model,
         expansion_ratio=expansion_ratio,
-        device=device,
+        device=device_str,
         bias=not no_bias,
         ffn_kwargs=ffn_config,
     )
@@ -59,14 +60,14 @@ def num_params(model: nn.Module) -> int:
     ffn2_numparams = num_params(ffn2)
     assert (
         ffn1_numparams == ffn2_numparams
-    ), "Only activation paths should have changed, re-check modeling!"
+    ), 'Only activation paths should have changed, re-check modeling!'
 
     input_ = torch.rand(1, d_model, device=device)
     output1 = ffn1(input_)
     output2 = ffn2(input_)
     assert (
         output1.numel() == output2.numel()
-    ), "Only activation paths should have changed, re-check modeling!"
+    ), 'Only activation paths should have changed, re-check modeling!'
     assert (
         not torch.allclose(output1, output2)
-    ), "Functions are different, outputs should not match!"
+    ), 'Functions are different, outputs should not match!'