fix: return averaged flops

marsggbo · Apr 29, 2022 · e1084b9 · e1084b9
1 parent be59c7b
commit e1084b9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/hyperbox/utils/calc_model_size.py b/hyperbox/utils/calc_model_size.py
@@ -146,7 +146,7 @@ def remove_buffer(m_):
         total_params += m.total_params
         # print("%s: %.2f %.2f" % (name, m.total_ops.item(), m.total_params.item()))
 
-    total_ops = total_ops.item()
+    total_ops = total_ops.item() / input_size[0]
     total_params = total_params.item()
 
     model.train(training).to(original_device)