bug fix

Author: warner-benjamin

optimi/optimizer.py

@@ -46,7 +46,7 @@ def __init__(self, params: Iterable[Tensor] | Iterable[dict], defaults: dict[str
 
         if not MIN_TORCH_2_6 and defaults.get("triton", False):
             raise ValueError(f"triton={defaults['triton']} requires PyTorch 2.6 or later. Set triton=False or upgrade PyTorch.")
-        elif not HAS_TRITON and defaults.get("triton", False):
+        if not HAS_TRITON and defaults.get("triton", False):
             raise ImportError("Triton could not be imported on this system. Set triton=False or install Triton.")
 
         super().__init__(params, defaults)

optimi/utils.py

@@ -19,6 +19,7 @@
 
 MIN_TORCH_2_1 = parse(torch.__version__) >= parse("2.1")
 MIN_TORCH_2_6 = parse(torch.__version__) >= parse("2.6")
+HAS_TRITON = HAS_TRITON and MIN_TORCH_2_6
 
 
 def debias(beta: float, step: int) -> float:
@@ -94,13 +95,9 @@ def _default_to_triton(params: list[torch.Tensor]) -> tuple[bool, bool]:
     if torch.jit.is_scripting():
         return False
 
-    triton = (
-        MIN_TORCH_2_6
-        and HAS_TRITON
-        and all(
-            p is None or (type(p) in _foreach_supported_types and _triton_kernels_supported_device(p) and torch.is_floating_point(p))
-            for p in params
-        )
+    triton = HAS_TRITON and all(
+        p is None or (type(p) in _foreach_supported_types and _triton_kernels_supported_device(p) and torch.is_floating_point(p))
+        for p in params
     )
     return triton