add info regarding FLOP calculation

tianyic · Jan 20, 2024 · cdc7c16 · cdc7c16
1 parent c47058d
commit cdc7c16
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/only_train_once/operation/operator.py b/only_train_once/operation/operator.py
@@ -116,6 +116,11 @@ def set_num_groups(self):
             self.num_groups = max(self.num_groups, param.shape[0])
 
     def compute_flops(self, input_shape):
+        '''
+        Currently only add FLOP calculation for conv and linear layers.
+        Other operator requires further implementations. 
+        Welcome PR in any.
+        '''
         return 0
 
     @property
@@ -127,6 +132,7 @@ def num_params(self):
         return num_params
 
 class ParamOTO(Operator):
+
     '''
     Operator for the tensor parameters in torch yet not formed in nn.Module 
     '''