fix nit

Signed-off-by: Stanley Winata <stanley.winata@amd.com>
iree-org · Sep 30, 2024 · 3d9db14 · 3d9db14
1 parent 7a13b1c
commit 3d9db14
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/shark_turbine/kernel/ops/wave_ops.py b/shark_turbine/kernel/ops/wave_ops.py
@@ -895,7 +895,7 @@ def captured_vars(self, graph: fx.Graph) -> list[fx.Node]:
         return captured_vars
 
     @property
-    def type(self) -> list[Memory | Register]:
+    def type(self) -> Memory | Register | list[Memory | Register]:
         res_types = [get_custom(x).type for x in self.init_args]
         if len(res_types) == 1:
             res_types = res_types[0]