From cdc7c1615f778bc0247c3914b4ca68012f351912 Mon Sep 17 00:00:00 2001
From: tianyi <ctydoctc@gmail.com>
Date: Sat, 20 Jan 2024 07:13:06 +0000
Subject: [PATCH] add info regarding FLOP calculation

---
 only_train_once/operation/operator.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/only_train_once/operation/operator.py b/only_train_once/operation/operator.py
index a6a5550..87fcec8 100644
--- 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 
     '''