diff --git a/generate_all_operators.py b/generate_all_operators.py
index 33580f31ff573e9c50f523f118c677b15789287f..2cae2eb47872d1d0f2e031fec29c3bc849f06b90 100644
--- a/generate_all_operators.py
+++ b/generate_all_operators.py
@@ -22,7 +22,7 @@ from functools import partial
 import logging
 import os
 import re
-from typing import Callable, List, Optional, Sequence, Set, Tuple
+from typing import Callable, List, Optional, Sequence, Set, Tuple, Union
 
 import numpy as np
 import sympy as sp
@@ -537,7 +537,7 @@ def all_operators(
     opts: List[Tuple[Set[Opts], LoopStrategy, str]],
     type_descriptor: HOGType,
     blending: GeometryMap,
-    geometries: Set[ElementGeometry],
+    geometries: Set[Union[TriangleElement, TetrahedronElement]],
 ) -> List[OperatorInfo]:
     P1 = LagrangianFunctionSpace(1, symbolizer)
     P1Vector = TensorialVectorFunctionSpace(P1)