type hint

PTB-MR · Nov 12, 2024 · 5c9f7a0 · 5c9f7a0
1 parent e23054b
commit 5c9f7a0
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/mrpro/algorithms/optimizers/pgd.py b/src/mrpro/algorithms/optimizers/pgd.py
@@ -4,12 +4,14 @@
 
 import torch
 
-from mrpro.operators.Functional import Functional, ProximableFunctional
+from mrpro.operators.Functional import ProximableFunctional
+from mrpro.operators.Operator import Operator
 
 
+# TODO: make it work with g:ProximableFunctionalSeparableSum and f:Operator with multiple inputs
 def pgd(
-    f: Functional,
-    g: ProximableFunctional,  # TODO: would it work with ProximableFunctionalSeparableSum?
+    f: Operator[torch.Tensor, tuple[torch.Tensor]],
+    g: ProximableFunctional,
     initial_value: torch.Tensor,
     stepsize: float = 1.0,
     reg_parameter: float = 0.01,