feat: Add compute FLOPs flag.

SysCV · Aug 7, 2024 · 456e3da · 456e3da
1 parent c834f28
commit 456e3da
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 1 deletion.
diff --git a/vis4d/common/ckpt.py b/vis4d/common/ckpt.py
@@ -205,7 +205,9 @@ def load_from_local(
     filename = osp.expanduser(filename)
     if not osp.isfile(filename):
         raise FileNotFoundError(f"{filename} can not be found.")
-    checkpoint = torch.load(filename, map_location=map_location)
+    checkpoint = torch.load(
+        filename, weights_only=True, map_location=map_location
+    )
     return checkpoint
 
 

diff --git a/vis4d/pl/run.py b/vis4d/pl/run.py
@@ -155,6 +155,7 @@ def main(argv: ArgsType) -> None:
         hyper_params,
         config.seed,
         ckpt_path if not resume else None,
+        config.compute_flops,
     )
     data_module = DataModule(config.data)
 

diff --git a/vis4d/zoo/base/runtime.py b/vis4d/zoo/base/runtime.py
@@ -58,6 +58,7 @@ def get_default_cfg(
     config.use_tf32 = False
     config.tf32_matmul_precision = "highest"
     config.benchmark = False
+    config.compute_flops = False
 
     return config