fix flop counter for trivial module

Summary: otherwise it fails.

Reviewed By: haooooooqi

Differential Revision: D29981975

fbshipit-source-id: dab1c8e404f4dfa7ec8eb373077009f435259dbf

fvcore/nn/print_model_statistics.py

@@ -155,7 +155,7 @@ def _remove_zero_statistics(
             with submodules removed if they have zero for all statistics.
     """
     out_stats: Dict[str, Dict[str, int]] = {}
-    _force_keep: Set[str] = force_keep if force_keep else set()
+    _force_keep: Set[str] = force_keep if force_keep else set() | {""}
 
     def keep_stat(name: str) -> None:
         prefix = name + ("." if name else "")

tests/test_print_model_statistics.py

@@ -535,3 +535,12 @@ def test_flop_count_str(self) -> None:
         # "    )\n"
         # "  )\n"
         # ")"
+
+    def test_flop_count_empty(self) -> None:
+        model = nn.ReLU()
+        inputs = (torch.randn((1, 10)),)
+        table = flop_count_table(FlopCountAnalysis(model, inputs))
+        self.assertGreater(len(table), 0)
+
+        out = flop_count_str(FlopCountAnalysis(model, inputs))
+        self.assertGreater(len(out), 0)