network forward can be handled externally

frame_averaging_pytorch/frame_averaging.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from random import randrange
 from optree import tree_flatten, tree_unflatten
 
@@ -19,7 +21,7 @@ def default(v, d):
 class FrameAverage(Module):
     def __init__(
         self,
-        net: Module,
+        net: Module | None = None,
         dim = 3,
         stochastic = False,
         invariant_output = False
@@ -146,7 +148,7 @@ def frame_average(out):
 
         # if one wants to handle the framed inputs externally
 
-        if return_framed_inputs_and_averaging_function:
+        if return_framed_inputs_and_averaging_function or not exists(self.net):
             return inputs, frame_average
 
         # merge frames into batch

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "frame-averaging-pytorch"
-version = "0.0.15"
+version = "0.0.16"
 description = "Frame Averaging"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }