[Operator] Fix code style of isclose and allclose op

FlagOpen · Jul 2, 2024 · 5762a06 · 5762a06
1 parent 96786cc
commit 5762a06
Showing 1 changed file with 1 addition and 7 deletions.
diff --git a/src/flag_gems/ops/isclose.py b/src/flag_gems/ops/isclose.py
@@ -92,13 +92,7 @@ def _isclose(
             "atol must be greater than or equal to zero, but got {}".format(atol)
         )
 
-    if (
-        A.dtype == torch.int64
-        or A.dtype == torch.int32
-        or A.dtype == torch.int16
-        or A.dtype == torch.int8
-        or A.dtype == torch.bool
-    ):
+    if A.dtype in (torch.int64, torch.int32, torch.int16, torch.int8, torch.bool):
         return isclose_func_int(A, B, rtol, atol)
     elif equal_nan:
         if A.dtype == torch.float32 or A.dtype == torch.float64: