Daily arc lint --take BLACK

Reviewed By: zertosh Differential Revision: D26481052 fbshipit-source-id: 21f102fa8bcfcec99dd4a0ff1b8def099a3aafae
xtinkt · Feb 17, 2021 · bf6401e · bf6401e
1 parent 9e53d28
commit bf6401e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/tests/test_mlp.py b/tests/test_mlp.py
@@ -11,13 +11,15 @@ class TestMLP(unittest.TestCase):
     and linear evaluation MLP layers
     """
 
-    MODEL_CONFIG = AttrDict({
-        "HEAD": {
-            "BATCHNORM_EPS": 1e-6,
-            "BATCHNORM_MOMENTUM": 0.99,
-            "PARAMS_MULTIPLIER": 1.0,
+    MODEL_CONFIG = AttrDict(
+        {
+            "HEAD": {
+                "BATCHNORM_EPS": 1e-6,
+                "BATCHNORM_MOMENTUM": 0.99,
+                "PARAMS_MULTIPLIER": 1.0,
+            }
         }
-    })
+    )
 
     def test_mlp(self):
         mlp = MLP(self.MODEL_CONFIG, dims=[2048, 100])

diff --git a/vissl/models/heads/mlp.py b/vissl/models/heads/mlp.py
@@ -92,8 +92,8 @@ def forward(self, batch: torch.Tensor):
             ), "MLP input should be either a tensor (2D, 4D) or list containing 1 tensor."
             batch = batch[0]
         if batch.ndim > 2:
-            assert (
-                    all(d == 1 for d in batch.shape[2:])
+            assert all(
+                d == 1 for d in batch.shape[2:]
             ), f"MLP expected 2D input tensor or 4D tensor of shape NxCx1x1. got: {batch.shape}"
             batch = batch.reshape((batch.size(0), batch.size(1)))
         out = self.clf(batch)