move actions from core, stop from being a provider module

wilsonmr · Apr 28, 2021 · c42a451 · c42a451
1 parent 209a337
commit c42a451
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 11 deletions.
diff --git a/anvil/core.py b/anvil/core.py
@@ -4,7 +4,6 @@
 import torch
 import torch.nn as nn
 
-from reportengine import collect
 
 ACTIVATION_LAYERS = {
     "leaky_relu": nn.LeakyReLU,
@@ -74,10 +73,3 @@ def forward(self, v_in: torch.tensor):
         shape (n_batch, size_out)
         """
         return self.network(v_in)
-
-_normalising_flow = collect("model_action", ("model_params",))
-
-def model_to_load(_normalising_flow):
-    return _normalising_flow[0]
-
-
diff --git a/anvil/models.py b/anvil/models.py
@@ -6,8 +6,9 @@
 """
 from functools import partial
 
-from anvil.core import Sequential
+from reportengine import collect
 
+from anvil.core import Sequential
 import anvil.layers as layers
 
 
@@ -103,6 +104,12 @@ def affine_spline(real_nvp, rational_quadratic_spline):
     return Sequential(real_nvp, rational_quadratic_spline)
 
 
+_normalising_flow = collect("model_action", ("model_params",))
+
+def model_to_load(_normalising_flow):
+    return _normalising_flow[0]
+
+
 MODEL_OPTIONS = {
     "nice": nice,
     "real_nvp": real_nvp,

diff --git a/anvil/scripts/anvil_sample.py b/anvil/scripts/anvil_sample.py
@@ -11,7 +11,6 @@
 log = logging.getLogger(__name__)
 
 PROVIDERS = [
-    "anvil.core",
     "anvil.models",
     "anvil.sample",
     "anvil.models",

diff --git a/anvil/scripts/anvil_train.py b/anvil/scripts/anvil_train.py
@@ -10,7 +10,7 @@
 
 log = logging.getLogger(__name__)
 
-PROVIDERS = ["anvil.train", "anvil.checkpoint", "anvil.core", "anvil.models"]
+PROVIDERS = ["anvil.train", "anvil.checkpoint", "anvil.models"]
 
 TRAINING_ACTIONS = ["train"]