fix ruff

ToluClassics · Jul 11, 2024 · fe4915f · fe4915f
1 parent f646a8e
commit fe4915f
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 4 deletions.
diff --git a/benchmark/run_benchmark.py b/benchmark/run_benchmark.py
@@ -18,7 +18,8 @@
     XLMRobertaModel,
 )
 
-import os, sys
+import os 
+import sys
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
 from src.mlx_transformers.models import BertForMaskedLM as MlxBertForMaskedLM

diff --git a/tests/test_llama.py b/tests/test_llama.py
@@ -28,4 +28,4 @@ def test_forward(self) -> None:
         inputs = {key: mx.array(v) for key, v in inputs.items()}
         outputs = self.model(**inputs, use_cache=True)
 
-        assert type(outputs.logits) == mx.array
+        assert type(outputs.logits) is mx.array
diff --git a/tests/test_phi.py b/tests/test_phi.py
@@ -28,4 +28,4 @@ def test_forward(self) -> None:
         inputs = {key: mx.array(v) for key, v in inputs.items()}
         outputs = self.model(**inputs, use_cache=True)
 
-        assert type(outputs.logits) == mx.array
+        assert type(outputs.logits) is mx.array
diff --git a/tests/test_phi3.py b/tests/test_phi3.py
@@ -32,4 +32,4 @@ def test_forward(self) -> None:
         inputs = {key: mx.array(v) for key, v in inputs.items()}
         outputs = self.model(**inputs, use_cache=True)
 
-        assert type(outputs.logits) == mx.array
+        assert type(outputs.logits) is mx.array