diff --git a/pypulseq/Sequence/write_seq.py b/pypulseq/Sequence/write_seq.py
index db5ab78..ce2aa3b 100644
--- a/pypulseq/Sequence/write_seq.py
+++ b/pypulseq/Sequence/write_seq.py
@@ -78,8 +78,8 @@ def write(self, file_name: str) -> None:
         output_file.write('\n')
 
     grad_lib_values = np.array(list(self.grad_library.type.values()))
-    arb_grad_mask = grad_lib_values == 'g'
-    trap_grad_mask = grad_lib_values == 't'
+    arb_grad_mask = [value == 'g' for value in grad_lib_values]     # Updated code
+    trap_grad_mask = [value == 't' for value in grad_lib_values]    # Updated code
 
     if any(arb_grad_mask):
         output_file.write('# Format of arbitrary gradients:\n')