use is_torch_version

Author: hlky

src/diffusers/models/activations.py

@@ -18,7 +18,7 @@
 from torch import nn
 
 from ..utils import deprecate
-from ..utils.import_utils import is_torch_npu_available
+from ..utils.import_utils import is_torch_npu_available, is_torch_version
 
 
 if is_torch_npu_available():
@@ -79,7 +79,7 @@ def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: b
         self.approximate = approximate
 
     def gelu(self, gate: torch.Tensor) -> torch.Tensor:
-        if gate.device.type == "mps" and torch.__version__ < '2.0.0':
+        if gate.device.type == "mps" and is_torch_version("<", "2.0.0"):
             # fp16 gelu not supported on mps before torch 2.0
             return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to(dtype=gate.dtype)
         return F.gelu(gate, approximate=self.approximate)
@@ -105,7 +105,7 @@ def __init__(self, dim_in: int, dim_out: int, bias: bool = True):
         self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias)
 
     def gelu(self, gate: torch.Tensor) -> torch.Tensor:
-        if gate.device.type == "mps" and torch.__version__ < '2.0.0':
+        if gate.device.type == "mps" and is_torch_version("<", "2.0.0"):
             # fp16 gelu not supported on mps before torch 2.0
             return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)
         return F.gelu(gate)