diff --git a/tanabe/__main__.py b/tanabe/__main__.py
index c1e6789..a6318fc 100644
--- a/tanabe/__main__.py
+++ b/tanabe/__main__.py
@@ -30,11 +30,11 @@ def __init__(
         self.delta_B = self.e_range / self.B
 
         self.spin_state = int(mode)
-        if self.spin_state in [4, 5, 6]:
+        if self.spin_state in {4, 5, 6}:
             self._size = 42
-        if self.spin_state in [3, 7]:
+        if self.spin_state in {3, 7}:
             self._size = 19
-        if self.spin_state in [2, 8]:
+        if self.spin_state in {2, 8}:
             self._size = 10
         self.result = np.zeros((self._size + 1, nroots))