chore: Applying lint

Signed-off-by: Anurag Dixit <[email protected]>

Author: andi4191

tests/modules/custom_models.py

@@ -3,6 +3,7 @@
 from transformers import BertModel, BertTokenizer, BertConfig
 import torch.nn.functional as F
 
+
 # Sample Pool Model (for testing plugin serialization)
 class Pool(nn.Module):
 
@@ -98,16 +99,15 @@ def BertModule():
     tokens_tensor = torch.tensor([indexed_tokens])
     segments_tensors = torch.tensor([segments_ids])
     config = BertConfig(
-                vocab_size_or_config_json_file=32000,
-                hidden_size=768,
-                num_hidden_layers=12,
-                num_attention_heads=12,
-                intermediate_size=3072,
-                torchscript=True,
-                )
+        vocab_size_or_config_json_file=32000,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        torchscript=True,
+    )
     model = BertModel(config)
     model.eval()
     model = BertModel.from_pretrained(model_name, torchscript=True)
     traced_model = torch.jit.trace(model, [tokens_tensor, segments_tensors])
     return traced_model
-

tests/modules/hub.py

@@ -111,24 +111,25 @@ def get(n, m, manifest):
     if n == "bert-base-uncased":
         traced_model = m["model"]
         torch.jit.save(traced_model, traced_filename)
-        manifest.update({n : [traced_filename]})
+        manifest.update({n: [traced_filename]})
     else:
         m["model"] = m["model"].eval().cuda()
         if m["path"] == "both" or m["path"] == "trace":
             trace_model = torch.jit.trace(m["model"], [x])
             torch.jit.save(trace_model, traced_filename)
-            manifest.update({n : [traced_filename]})
+            manifest.update({n: [traced_filename]})
         if m["path"] == "both" or m["path"] == "script":
             script_model = torch.jit.script(m["model"])
             torch.jit.save(script_model, script_filename)
             if n in manifest.keys():
                 files = list(manifest[n]) if type(manifest[n]) != list else manifest[n]
                 files.append(script_filename)
-                manifest.update({n : files})
+                manifest.update({n: files})
             else:
                 manifest.update({n: [script_filename]})
     return manifest
 
+
 def download_models(version_matches, manifest):
     # Download all models if torch version is different than model version
     if not version_matches:
@@ -142,8 +143,8 @@ def download_models(version_matches, manifest):
             if (m["path"] == "both" and os.path.exists(scripted_filename) and os.path.exists(traced_filename)) or \
                (m["path"] == "script" and os.path.exists(scripted_filename)) or \
                (m["path"] == "trace" and os.path.exists(traced_filename)):
-                   print("Skipping {} ".format(n))
-                   continue
+                print("Skipping {} ".format(n))
+                continue
             manifest = get(n, m, manifest)
 
 
@@ -184,4 +185,5 @@ def main():
         f.write(record)
         f.truncate()
 
-main()
+
+main()