add suggested __call__ function

lab-cosmo · Mar 17, 2023 · 727f9b0 · 727f9b0
1 parent 13fd083
commit 727f9b0
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/equisolve/module.py b/src/equisolve/module.py
@@ -17,9 +17,11 @@ class Module(metaclass=ABCMeta):
         def forward(self, *args, **kwargs):
             return
 
+        def __call__(self, *args, **kwargs):
+            self.forward(*args, **kwargs)
 
-TEstimatorModule = TypeVar("TEstimatorModule", bound="EstimatorModule")
 
+TEstimatorModule = TypeVar("TEstimatorModule", bound="EstimatorModule")
 
 class EstimatorModule(Module, metaclass=ABCMeta):
     def forward(self, X: TensorMap):