diff --git a/reasoning_gym/cognition/color_cube_rotation.py b/reasoning_gym/cognition/color_cube_rotation.py
index ef2210ee..b0b65fae 100644
--- a/reasoning_gym/cognition/color_cube_rotation.py
+++ b/reasoning_gym/cognition/color_cube_rotation.py
@@ -119,7 +119,7 @@ class ColorCubeRotationDataset(ProceduralDataset):
         # Generate sequence of rotations
         num_rotations = rng.randint(self.config.min_rotations, self.config.max_rotations)
         rotations = []
-        
+
         # Keep trying until we have at least one valid rotation
         while len(rotations) < num_rotations:
             # Get all sides except TOP
diff --git a/reasoning_gym/cognition/number_sequences.py b/reasoning_gym/cognition/number_sequences.py
index c674ed49..bc5770df 100644
--- a/reasoning_gym/cognition/number_sequences.py
+++ b/reasoning_gym/cognition/number_sequences.py
@@ -1,12 +1,12 @@
 from dataclasses import dataclass
-from enum import Enum
+from enum import StrEnum
 from random import Random
 from typing import List, Optional
 
 from ..dataset import ProceduralDataset
 
 
-class Operation(Enum):
+class Operation(StrEnum):
     """Basic mathematical operations that can be composed"""
 
     ADD = "+"
diff --git a/reasoning_gym/logic/propositional_logic.py b/reasoning_gym/logic/propositional_logic.py
index 8897d2fa..565527ad 100644
--- a/reasoning_gym/logic/propositional_logic.py
+++ b/reasoning_gym/logic/propositional_logic.py
@@ -1,12 +1,12 @@
 """Propositional logic task generator"""
 
 from dataclasses import dataclass
-from enum import Enum
+from enum import StrEnum
 from random import Random
-from typing import Any, List, Optional, Set, Tuple
+from typing import Any, List, Optional, Set
 
 
-class Operator(Enum):
+class Operator(StrEnum):
     """Basic logical operators"""
 
     AND = "∧"