Refactor amp runner

MauroCamporeale · Apr 13, 2021 · 02baa30 · 02baa30
1 parent d2ee4c8
commit 02baa30
Show file tree

Hide file tree

Showing 5 changed files with 61 additions and 974 deletions.
diff --git a/mmcv_custom/runner/__init__.py b/mmcv_custom/runner/__init__.py
@@ -1,11 +1,8 @@
 # Copyright (c) Open-MMLab. All rights reserved.
-from .base_runner import BaseRunner
-from .checkpoint import (_load_checkpoint, load_checkpoint, load_state_dict,
-                         save_checkpoint, weights_to_cpu)
+from .checkpoint import save_checkpoint
 from .epoch_based_runner import EpochBasedRunnerAmp
 
 
 __all__ = [
-    'BaseRunner', 'EpochBasedRunnerAmp', '_load_checkpoint', 'load_checkpoint',
-    'load_state_dict', 'save_checkpoint', 'weights_to_cpu'
+    'EpochBasedRunnerAmp', 'save_checkpoint'
 ]